Goの勉強として、TCPプロキシを作りました。
任意のTCP接続の間にはさまって、通信ログを取ったりできます。
このようなイメージです。
構成要素
- TCP接続を受け入れる
- 受け入れ後、本物のサーバーにつなぐ
- その後は 受け入れたSocketと本物のサーバーを双方向にリレーする
TCP接続受け入れ
まずは外部からのTCP接続を受け入れる部分を作ります。
普通にTCPサーバーを作るときと同様です。
新しい接続がきたら handleConn 関数を goroutine で実行することで複数の接続を受け入れ可能にしています。
handleConn の中身はこの後に記述していきます。
type Server struct {}
func (s *Server) Start(listenAddr *net.TCPAddr) error {
lt, err := net.ListenTCP("tcp", listenAddr)
if err != nil {
return err
}
defer lt.Close()
for {
conn, err := lt.AcceptTCP()
if err != nil {
return err
}
go s.handleConn(conn)
}
}
func (s *Server) handleConn(c *net.TCPConn) {
...
}
受け入れ後の双方向リレー
クライアントから接続したあとは、本物のサーバーへ接続して、
クライアントとサーバーの通信内容をリレー（双方向につなげる）します。
双方向ということは、
- クライアント→サーバー
- サーバー→クライアント
の2方向なので、2つの goroutine を並列させて行います。
func (s *Server) handleConn(c *net.TCPConn) {
defer c.Close()
// 何らかの方法で本物のサーバーのアドレスを取得しておく
serverAddr := ....
// 本物のサーバーへ接続
serverConn, err := net.DialTCP("tcp", nil, serverAddr)
if err != nil {
return err
}
defer serverConn.Close()
// 双方向のリレーを作る
p := &Proxy{}
p.Start(c, serverConn)
}
// プロキシ本体
type Proxy struct{}
// プロキシの開始（双方向リレー開始）
func (p *Proxy) Start(clientConn, serverConn) error {
defer clientConn.Close()
defer serverConn.Close()
var eg errGroup.Group
eg.Go(func() error { return p.relay(&eg, clientConn, serverConn) })
eg.Go(func() error { return p.relay(&eg, serverConn, clientConn) })
// wait for stop
return eg.Wait()
}
const (
BUFFER_SIZE = 0xFFFF
)
// fromConn -> toConn へ通信内容をリレーする
func (p *Proxy) relay(eg *errGroup.Group, fromConn, toConn *net.TCPConn) error {
buff := make([]byte, BUFFER_SIZE)
for {
n, err := fromConn.Read(buff)
if err != nil {
return err
}
b := buff[:n]
// ここでログなどを取ることができる
n, err = toConn.Write(b)
if err != nil {
return err
}
}
}
まとめ
- サーバー部分とリレー部分で分けて考えるとわかりやすい
- goroutineをまたぐ通知は channel をつかう
- 複数のgoroutineの待機をしたい場合は sync.WaitGroup や errGroup.Group を使うと良い（キャンセルなどもしたい場合は、context を使う？）