网络代理【6】TCP代理，代理Redis服务

今天药忘吃喽~ 2022-10-21 11:51 134阅读 0赞

> 前置文章:[【TCP服务器】仿http.server源码实现服务与逻辑分离][TCP_http.server]  
> 本文是在以上文章的TCP Handler接口实现

# 一、代理实现Handler接口 #

package proxy
    
    // TCPReverseProxy TCP反向代理
    type TCPReverseProxy struct { 
    	ctx                  context.Context //单次请求单独设置
    	Addr                 string
    	KeepAlivePeriod      time.Duration //设置
    	DialTimeout          time.Duration //设置超时时间
    	DialContext          func(ctx context.Context, network, address string) (net.Conn, error)
    	OnDialError          func(src net.Conn, dstDialErr error)
    	ProxyProtocolVersion int
    }
    
    func (s *TCPReverseProxy) dialTimeout() time.Duration { 
    	if s.DialTimeout > 0 { 
    		return s.DialTimeout
    	}
    	return 10 * time.Second
    }
    
    func (s *TCPReverseProxy) keepAlivePeriod() time.Duration { 
    	if s.KeepAlivePeriod != 0 { 
    		return s.KeepAlivePeriod
    	}
    	return time.Minute
    }
    
    func (s *TCPReverseProxy) dialContext() func(ctx context.Context, network, address string) (net.Conn, error) { 
    	if s.dialContext() != nil { 
    		return s.DialContext
    	}
    	return (&net.Dialer{ 
    		Timeout:   s.DialTimeout,
    		KeepAlive: s.KeepAlivePeriod,
    	}).DialContext
    }
    
    func (s *TCPReverseProxy) onDialError() func(src net.Conn, dstDialErr error) { 
    	if s.OnDialError != nil { 
    		return s.OnDialError
    	}
    	return func(src net.Conn, dstDialErr error) { 
    		log.Printf("tcpproxy: for incoming conn %v, error dialing %q: %v", src.RemoteAddr().String(), s.Addr, dstDialErr)
    		src.Close()
    	}
    }
    
    func (s *TCPReverseProxy) proxyCopy(errc chan<- error, dst, src net.Conn) { 
    	_, err := io.Copy(dst, src)
    	errc <- err
    }
    
    // ServeTCP 传入上游conn, 在这里完成下游链接与数据交换
    func (s *TCPReverseProxy) ServeTCP(ctx context.Context, src net.Conn) { 
    	// 设置连接超时
    	var cancel context.CancelFunc
    	if s.DialTimeout >= 0 { 
    		ctx, cancel = context.WithTimeout(ctx, s.dialTimeout())
    	}
    	dst, err := s.dialContext()(ctx, "tcp", s.Addr)
    	if cancel != nil { 
    		cancel()
    	}
    	if err != nil { 
    		s.onDialError()(src, err)
    		return
    	}
    	defer func() {  // 退出下游连接
    		go dst.Close()
    	}()
    
    	// 设置dst的keepalive参数，在数据请求之前
    	if ka := s.keepAlivePeriod(); ka > 0 { 
    		if c, ok := dst.(*net.TCPConn); ok { 
    			c.SetKeepAlive(true)
    			c.SetKeepAlivePeriod(ka)
    		}
    	}
    	errc := make(chan error, 1)
    	go s.proxyCopy(errc, src, dst)
    	go s.proxyCopy(errc, dst, src)
    	<-errc
    }
    
    func NewMultiTCPReverseProxy(addr string) *TCPReverseProxy { 
    	return &TCPReverseProxy{ 
    		ctx:             context.Background(),
    		Addr:            addr,
    		KeepAlivePeriod: time.Second,
    		DialTimeout:     time.Second,
    	}
    }

# 二、启动Redis服务 #

docker run -itd --name redis -p 6379:6379 redis

# 三、创建Redis TCP代理并测试 #

package main
    
    import (
    	"demo/proxy"
    	"demo/server/tcp"
    )
    
    type handler struct{ }
    
    func (t *handler) ServeTCP(ctx context.Context, conn net.Conn) { 
    	conn.Write([]byte("tcp handler\n"))
    }
    
    var addr = ":3001"
    
    func main() { 
    	log.Println("start tcp-server at ", addr)
    	proxy := proxy.NewMultiTCPReverseProxy("127.0.0.1:6379")
    	srv := tcp.Server{ Addr: addr, Handler: proxy}
    	log.Fatal(srv.ListenAndServe())
    }

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM4OTAwNTY1_size_16_color_FFFFFF_t_70_pic_center]

[TCP_http.server]: https://blog.csdn.net/qq_38900565/article/details/116379958
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM4OTAwNTY1_size_16_color_FFFFFF_t_70_pic_center]: /images/20221021/7099bab40cd046d39788bbcb4236ed84.png