go 实现TCP代理

TCP服务器

参考 go创建tcp服务

package main

import (
	"context"
	"fmt"
	"net"
	"runtime"
	"time"
)

//优化了下，定义Handler接口
type Handler interface {
	TcpServer(c context.Context, con net.Conn)
}

type TcpService struct {
	//地址
	Addr string
	//增加了上下文
	Handler Handler
	//初始上下文
	Context context.Context

	//增加读取超时 写入超时 连接探活超时
	ReadTimeOut      time.Duration
	WriteTimeOut     time.Duration
	KeepAliveTimeOut time.Duration
}

//关闭Conn
type Conn struct {
	service *TcpService
	conn    net.Conn
}

func (conn Conn) Server(c context.Context) {

	//错误堆栈信息输出
	defer func() {
		if err := recover(); err != nil {
			const size = 64 << 10
			buf := make([]byte, size)
			//runtime.Stack 堆栈写入buf,并返回写入长度
			buf = buf[:runtime.Stack(buf, false)]
			//输出错误堆栈
			fmt.Println(string(buf))
		}
		conn.Close()
	}()
	conn.service.Handler.TcpServer(c, conn.conn)
}

func (conn Conn) Close() {
	conn.conn.Close()
}

//封装成自己的连接
func (t *TcpService) NewCon(con net.Conn) *Conn {
	c := &Conn{
		service: t,
		conn:    con,
	}

	if d := t.ReadTimeOut; d != 0 {
		c.conn.SetReadDeadline(time.Now().Add(d))
	}

	if d := t.WriteTimeOut; d != 0 {
		c.conn.SetWriteDeadline(time.Now().Add(d))
	}

	if d := t.KeepAliveTimeOut; d != 0 {
		if tcpConn, ok := c.conn.(*net.TCPConn); ok {
			tcpConn.SetKeepAlive(true)
			tcpConn.SetKeepAlivePeriod(d)
		}
	}

	return c
}

//向http一样提供 ListenAndServe 启动服务器
func (t *TcpService) ListenAndServe() error {
	listen, err := net.Listen("tcp", t.Addr)
	if err != nil {
		panic(err)
	}

	defer listen.Close()

	//tcpListener := listen.(*net.TCPListener)
	//tcp, err := tcpListener.AcceptTCP()
	if err != nil {
		panic(err)
	}

	listener := listen.(*net.TCPListener)

	for {
		accept, err := listener.AcceptTCP()

		if err != nil {
			continue
		}

		c := t.NewCon(accept)

		go c.Server(t.Context)
	}
}

type DefaultTcpServer struct {
}

func (d DefaultTcpServer) TcpServer(c context.Context, con net.Conn) {
	con.Write([]byte("this is tcp handler"))
}

func main() {
	service := TcpService{
		Addr:    ":9001",
		Handler: DefaultTcpServer{},
	}
	service.ListenAndServe()
}

代理原理图

没有代理情况下的客户端与服务端直连

代理下的客户端，tcp服务器，代理服务器

代理服务

代理服务的部分代码时基于tcp服务器的,使用tcp服务代理tcp服务
代理的主要逻辑就是 使用 io.copy 进行上下net.conn 数据的交换

package main

import (
	"context"
	"fmt"
	"io"
	"net"
	"runtime"
	"time"
)

//优化了下，定义Handler接口
type Handler interface {
	TcpServer(c context.Context, con net.Conn)
}

type TcpService struct {
	//地址
	Addr string
	//增加了上下文
	Handler Handler
	//初始上下文
	Context context.Context

	//增加读取超时 写入超时 连接探活超时
	ReadTimeOut      time.Duration
	WriteTimeOut     time.Duration
	KeepAliveTimeOut time.Duration
}

//关闭Conn
type Conn struct {
	service *TcpService
	conn    net.Conn
}

func (conn Conn) Server(c context.Context) {

	//错误堆栈信息输出
	defer func() {
		if err := recover(); err != nil {
			const size = 64 << 10
			buf := make([]byte, size)
			//runtime.Stack 堆栈写入buf,并返回写入长度
			buf = buf[:runtime.Stack(buf, false)]
			//输出错误堆栈
			fmt.Println(string(buf))
		}
		conn.Close()
	}()
	conn.service.Handler.TcpServer(c, conn.conn)
}

func (conn Conn) Close() {
	conn.conn.Close()
}

//封装成自己的连接
func (t *TcpService) NewCon(con net.Conn) *Conn {
	c := &Conn{
		service: t,
		conn:    con,
	}

	if d := t.ReadTimeOut; d != 0 {
		c.conn.SetReadDeadline(time.Now().Add(d))
	}

	if d := t.WriteTimeOut; d != 0 {
		c.conn.SetWriteDeadline(time.Now().Add(d))
	}

	if d := t.KeepAliveTimeOut; d != 0 {
		if tcpConn, ok := c.conn.(*net.TCPConn); ok {
			tcpConn.SetKeepAlive(true)
			tcpConn.SetKeepAlivePeriod(d)
		}
	}

	return c
}

//向http一样提供 ListenAndServe 启动服务器
func (t *TcpService) ListenAndServe() error {
	listen, err := net.Listen("tcp", t.Addr)
	if err != nil {
		panic(err)
	}

	defer listen.Close()

	//tcpListener := listen.(*net.TCPListener)
	//tcp, err := tcpListener.AcceptTCP()
	if err != nil {
		panic(err)
	}

	listener := listen.(*net.TCPListener)

	for {
		accept, err := listener.AcceptTCP()

		if err != nil {
			continue
		}

		c := t.NewCon(accept)

		go c.Server(t.Context)
	}
}

type TcpReverseProxy struct {
	Addr string

	//超时设置
	TimeOut   time.Duration
	KeepAlive time.Duration
}

func (t *TcpReverseProxy) TcpServer(c context.Context, src net.Conn) {
	dst, err := t.DialContext()(context.Background(), "tcp", t.Addr)
	defer func() {
		src.Close()
		dst.Close()
	}()
	if err != nil {
		fmt.Printf("dail error %v", err)
		return
	}

	e := make(chan error, 1)
	go t.CopyTo(e, dst, src)
	go t.CopyTo(e, src, dst)
	<-e
}

//由于要设置连接超时时间和探活时间，使用net.Dailer.dailContext() 进行连接
func (t *TcpReverseProxy) DialContext() func(ctx context.Context, network, address string) (net.Conn, error) {
	return (&net.Dialer{
		Timeout:   t.TimeOut,
		KeepAlive: t.KeepAlive,
	}).DialContext
}

func (t *TcpReverseProxy) CopyTo(e chan<- error, dst, src net.Conn) {
	_, err := io.Copy(dst, src)
	e <- err
}

func main() {
	t := TcpService{
		Addr: ":9002",
		Handler: &TcpReverseProxy{
			Addr:      ":9001",
			TimeOut:   10 * time.Second,
			KeepAlive: 10 * time.Second,
		},
	}

	t.ListenAndServe()
}

代理redis

将代理地址更换成redis的地址

func main() {
	t := TcpService{
		Addr: ":9002",
		Handler: &TcpReverseProxy{
			Addr:      ":6379",
			TimeOut:   10 * time.Second,
			KeepAlive: 10 * time.Second,
		},
	}

	t.ListenAndServe()
}

通过 telnet 进行测试

xieruixiang@xieruixiangdeMacBook-Pro ~ % telnet 127.0.0.1 9002
Trying 127.0.0.1...
Connected to localhost.
Escape character is '^]'.
set a 2
+OK
get a
$1
2