feat: Add TLS support to gnet

acceptor.go

@@ -50,6 +50,11 @@ func (eng *engine) accept(fd int, _ netpoll.IOEvent) error {
 
 	el := eng.lb.next(remoteAddr)
 	c := newTCPConn(nfd, el, sa, el.ln.addr, remoteAddr)
+	if el.engine.opts.TLSconfig != nil {
+		if err = c.UpgradeTLS(el.engine.opts.TLSconfig); err != nil {
+			return err
+		}
+	}
 
 	err = el.poller.UrgentTrigger(el.register, c)
 	if err != nil {
@@ -84,6 +89,11 @@ func (el *eventloop) accept(fd int, ev netpoll.IOEvent) error {
 	}
 
 	c := newTCPConn(nfd, el, sa, el.ln.addr, remoteAddr)
+	if el.engine.opts.TLSconfig != nil {
+		if err = c.UpgradeTLS(el.engine.opts.TLSconfig); err != nil {
+			return err
+		}
+	}
 	if err = el.poller.AddRead(c.pollAttachment); err != nil {
 		return err
 	}

connection.go

@@ -33,6 +33,7 @@ import (
 	"github.com/panjf2000/gnet/v2/pkg/buffer/elastic"
 	gerrors "github.com/panjf2000/gnet/v2/pkg/errors"
 	bsPool "github.com/panjf2000/gnet/v2/pkg/pool/byteslice"
+	"github.com/panjf2000/gnet/v2/pkg/tls"
 )
 
 type conn struct {
@@ -48,6 +49,7 @@ type conn struct {
 	fd             int                     // file descriptor
 	isDatagram     bool                    // UDP protocol
 	opened         bool                    // connection opened event fired
+	tlsconn        *tls.Conn               // tls connection
 }
 
 func newTCPConn(fd int, el *eventloop, sa unix.Sockaddr, localAddr, remoteAddr net.Addr) (c *conn) {
@@ -127,8 +129,18 @@ func (c *conn) open(buf []byte) error {
 	return err
 }
 
+func (c *conn) writeTLS(data []byte) (n int, err error) {
+	// use tls to encrypt the data before sending it.
+	// tlsconn will call gnet.WriteTCP() to sent the data directly.
+	// If gnetConn.outboundBufferis not empty, data will be
+	// buffered in gnetConn.outboundBuffer.
+	n, err = c.tlsconn.Write(data)
+	return
+}
+
 func (c *conn) write(data []byte) (n int, err error) {
 	n = len(data)
+
 	// If there is pending data in outbound buffer, the current data ought to be appended to the outbound buffer
 	// for maintaining the sequence of network packets.
 	if !c.outboundBuffer.IsEmpty() {
@@ -154,6 +166,28 @@ func (c *conn) write(data []byte) (n int, err error) {
 	return
 }
 
+func (c *conn) writevTLS(bs [][]byte) (n int, err error) {
+	for _, b := range bs {
+		n += len(b)
+	}
+
+	// use tls to encrypt the data before sending it.
+	// tlsconn will call gnet.WriteTCP() to sent the data directly.
+	// If gnetConn.outboundBufferis not empty, data will be
+	// buffered in gnetConn.outboundBuffer.
+	sent := 0
+	var sentN int
+	for _, b := range bs {
+		sentN, err = c.tlsconn.Write(b)
+		if sentN < 0 {
+			// the connection is closed (c.loop.closeConn() is called).
+			return sent, err
+		}
+		sent += sentN
+	}
+	return
+}
+
 func (c *conn) writev(bs [][]byte) (n int, err error) {
 	for _, b := range bs {
 		n += len(b)
@@ -205,7 +239,11 @@ func (c *conn) asyncWrite(itf interface{}) (err error) {
 	}
 
 	hook := itf.(*asyncWriteHook)
-	_, err = c.write(hook.data)
+	if c.tlsconn != nil {
+		_, err = c.writeTLS(hook.data)
+	} else {
+		_, err = c.write(hook.data)
+	}
 	if hook.callback != nil {
 		_ = hook.callback(c, err)
 	}
@@ -223,7 +261,11 @@ func (c *conn) asyncWritev(itf interface{}) (err error) {
 	}
 
 	hook := itf.(*asyncWritevHook)
-	_, err = c.writev(hook.data)
+	if c.tlsconn != nil {
+		_, err = c.writevTLS(hook.data)
+	} else {
+		_, err = c.writev(hook.data)
+	}
 	if hook.callback != nil {
 		_ = hook.callback(c, err)
 	}
@@ -348,13 +390,25 @@ func (c *conn) Write(p []byte) (int, error) {
 		}
 		return len(p), nil
 	}
+	if c.tlsconn != nil {
+		return c.writeTLS(p)
+	}
+	return c.write(p)
+}
+
+// Expose the plaintext write API which should only be used
+// by tlsconn.Write().
+func (c *conn) WriteTCP(p []byte) (int, error) {
 	return c.write(p)
 }
 
 func (c *conn) Writev(bs [][]byte) (int, error) {
 	if c.isDatagram {
 		return 0, gerrors.ErrUnsupportedOp
 	}
+	if c.tlsconn != nil {
+		return c.writevTLS(bs)
+	}
 	return c.writev(bs)
 }
 
@@ -467,3 +521,27 @@ func (c *conn) Close() error {
 		return
 	}, nil)
 }
+
+func (c *conn) UpgradeTLS(config *tls.Config) (err error) {
+	// TODO: create a sync.pool to manage the TLS connection
+	c.tlsconn = tls.Server(c, config.Clone())
+
+	// 很有可能握手包在UpgradeTls之前发过来了，这里把inboundBuffer剩余数据当做握手数据处理
+	if c.inboundBuffer.Len() > 0 {
+		head, tail := c.inboundBuffer.Peek(-1)
+		c.tlsconn.RawInputSet(head) //nolint:errcheck
+		c.tlsconn.RawInputSet(tail) //nolint:errcheck
+		c.inboundBuffer.Reset()
+		if err := c.tlsconn.Handshake(); err != nil {
+			return err
+		}
+	}
+
+	// 握手失败的关了
+	time.AfterFunc(time.Second*5, func() {
+		if c.opened && (c.tlsconn == nil || !c.tlsconn.HandshakeComplete()) {
+			c.Close()
+		}
+	})
+	return err
+}

eventloop.go

@@ -114,7 +114,58 @@ func (el *eventloop) open(c *conn) error {
 	return el.handleAction(c, action)
 }
 
+func (el *eventloop) readTLS(c *conn) error {
+	// Since the el.Buffer may contain multiple TLS record,
+	// we process one TLS record in each iteration until no more
+	// TLS records are available
+	for {
+		if err := c.tlsconn.ReadFrame(); err != nil {
+			// Receive error unix.EAGAIN, wait for the next round
+			if err == unix.EAGAIN {
+				c.tlsconn.DataDone()
+				return nil
+			}
+			// If err is io.EOF, it can either the data is drained,
+			// receives a close notify from the client.
+			return el.closeConn(c, os.NewSyscallError("TLS read", err))
+		}
+
+		// load all decrypted data and make it ready for gnet to use
+		c.buffer = c.tlsconn.Data()
+
+		action := el.eventHandler.OnTraffic(c)
+		switch action {
+		case None:
+		case Close:
+			// tls data will be cleaned up in el.closeConn()
+			return el.closeConn(c, nil)
+		case Shutdown:
+			c.tlsconn.DataDone()
+			return gerrors.ErrEngineShutdown
+		}
+		_, _ = c.inboundBuffer.Write(c.buffer)
+		c.buffer = c.buffer[:0]
+
+		// all available TLS records are processed
+		if !c.tlsconn.IsRecordCompleted(c.tlsconn.RawInputData()) {
+			c.tlsconn.DataDone()
+			return nil
+		}
+	}
+}
+
 func (el *eventloop) read(c *conn) error {
+	// detected whether kernel TLS RX is enabled
+	// This only happens after TLS handshake is completed.
+	// Therefore, no need to call c.tlsconn.HandshakeComplete().
+	if c.tlsconn != nil && c.tlsconn.IsKTLSRXEnabled() {
+		// attach the gnet eventloop.buffer to tlsconn.rawInput.
+		// So, KTLS can decrypt the data directly to the buffer without memory allocation.
+		// Since data is read through KTLS, there is no need to call unix.read(c.fd, el.buffer)
+		c.tlsconn.RawInputSet(el.buffer) //nolint:errcheck
+		return el.readTLS(c)
+	}
+
 	n, err := unix.Read(c.fd, el.buffer)
 	if err != nil || n == 0 {
 		if err == unix.EAGAIN {
@@ -126,6 +177,29 @@ func (el *eventloop) read(c *conn) error {
 		return el.closeConn(c, os.NewSyscallError("read", err))
 	}
 
+	if c.tlsconn != nil {
+		// attach the gnet eventloop.buffer to tlsconn.rawInput.
+		c.tlsconn.RawInputSet(el.buffer[:n]) //nolint:errcheck
+		if !c.tlsconn.HandshakeComplete() {
+			// 先判断是否足够一条消息
+			data := c.tlsconn.RawInputData()
+			if !c.tlsconn.IsRecordCompleted(data) {
+				c.tlsconn.DataDone()
+				return nil
+			}
+			if err = c.tlsconn.Handshake(); err != nil {
+				// closeConn will cleanup the TLS data at the end,
+				// so no need to call tlsconn.DataDone()
+				return el.closeConn(c, os.NewSyscallError("TLS handshake", err))
+			}
+			if !c.tlsconn.HandshakeComplete() || len(c.tlsconn.RawInputData()) == 0 { // 握手没成功，或者握手成功，但是没有数据黏包了
+				c.tlsconn.DataDone()
+				return nil
+			}
+		}
+		return el.readTLS(c)
+	}
+
 	c.buffer = el.buffer[:n]
 	action := el.eventHandler.OnTraffic(c)
 	switch action {
@@ -192,6 +266,16 @@ func (el *eventloop) closeConn(c *conn, err error) (rerr error) {
 		return
 	}
 
+	// close the TLS connection by sending the alert
+	if c.tlsconn != nil {
+		// close the TLS connection, which will send a close notify to the client
+		c.tlsconn.Close()
+		// Make sure all memory requested from the pool is returned.
+		c.tlsconn.DataCleanUpAfterClose()
+		c.tlsconn = nil
+		// TODO: create a sync.pool to manage the TLS connection
+	}
+
 	// Send residual data in buffer back to the peer before actually closing the connection.
 	if !c.outboundBuffer.IsEmpty() {
 		for !c.outboundBuffer.IsEmpty() {

go.mod

@@ -1,11 +1,12 @@
 module github.com/panjf2000/gnet/v2
 
 require (
+	github.com/0-haha/gnet_go_tls/v120 v120.2.0
 	github.com/panjf2000/ants/v2 v2.7.1
-	github.com/stretchr/testify v1.8.1
+	github.com/stretchr/testify v1.8.2
 	github.com/valyala/bytebufferpool v1.0.0
 	go.uber.org/zap v1.21.0
-	golang.org/x/sys v0.3.0
+	golang.org/x/sys v0.4.0
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 )
 
@@ -15,7 +16,8 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	go.uber.org/atomic v1.10.0 // indirect
 	go.uber.org/multierr v1.8.0 // indirect
+	golang.org/x/crypto v0.5.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
-go 1.17
+go 1.20

go.sum

@@ -1,3 +1,5 @@
+github.com/0-haha/gnet_go_tls/v120 v120.2.0 h1:uPq6/VoULSna7DHMR3KELdALSUaNxZ0tjFts9Yneiso=
+github.com/0-haha/gnet_go_tls/v120 v120.2.0/go.mod h1:ZDwYfvBBzRwvZNENOXeVI+QjVDKr5r8aEKLMDSchRkI=
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
@@ -23,8 +25,9 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
+github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
@@ -40,6 +43,8 @@ go.uber.org/zap v1.21.0 h1:WefMeulhovoZ2sYXz7st6K0sLj7bBhpiFaud4r4zST8=
 go.uber.org/zap v1.21.0/go.mod h1:wjWOCqI0f2ZZrJF/UufIOkiC8ii6tm1iqIsLo76RfJw=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
+golang.org/x/crypto v0.5.0 h1:U/0M97KRkSFvyD/3FSmdP5W5swImpNgle/EHFhOsQPE=
+golang.org/x/crypto v0.5.0/go.mod h1:NK/OQwhpMQP3MwtdjgLlYHnH9ebylxKWv3e0fK+mkQU=
 golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
@@ -55,8 +60,8 @@ golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.3.0 h1:w8ZOecv6NaNa/zC8944JTU3vz4u6Lagfk4RPQxv92NQ=
-golang.org/x/sys v0.3.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.4.0 h1:Zr2JFtRQNX3BCZ8YtxRE9hNJYC8J6I1MVbMg6owUp18=
+golang.org/x/sys v0.4.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=

options.go

@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"github.com/panjf2000/gnet/v2/pkg/logging"
+	"github.com/panjf2000/gnet/v2/pkg/tls"
 )
 
 // Option is a function that will set up option.
@@ -109,6 +110,9 @@ type Options struct {
 	// SocketSendBuffer sets the maximum socket send buffer in bytes.
 	SocketSendBuffer int
 
+	// TLSconfig sets the configuration of a TLS connection
+	TLSconfig *tls.Config
+
 	// LogPath the local path where logs will be written, this is the easiest way to set up logging,
 	// gnet instantiates a default uber-go/zap logger with this given log path, you are also allowed to employ
 	// you own logger during the lifetime by implementing the following log.Logger interface.
@@ -249,3 +253,10 @@ func WithMulticastInterfaceIndex(idx int) Option {
 		opts.MulticastInterfaceIndex = idx
 	}
 }
+
+// WithTLS sets the tls configuration which includes the cert, the key, the cipher suite, the rotocol version, and etc.
+func WithTLS(tlsconfig *tls.Config) Option {
+	return func(opts *Options) {
+		opts.TLSconfig = tlsconfig
+	}
+}