kcp: introduce rekey support

Author: lonnywong

tsshd/bus.go

@@ -129,6 +129,8 @@ func handleBusEvent(stream Stream) {
 			err = handleAliveEvent(stream, heartbeatTimeout, intervalTime)
 		case "setting":
 			err = handleSettingEvent(stream)
+		case "rekey":
+			err = handleRekeyEvent(stream)
 		default:
 			if err := handleUnknownEvent(stream, command); err != nil {
 				warning("handle bus command [%s] failed: %v. You may need to upgrade tsshd.", command, err)
@@ -220,6 +222,14 @@ func handleSettingEvent(stream Stream) error {
 	return nil
 }
 
+func handleRekeyEvent(stream Stream) error {
+	var msg rekeyMessage
+	if err := recvMessage(stream, &msg); err != nil {
+		return fmt.Errorf("recv rekey message failed: %v", err)
+	}
+	return globalProtoServer.handleRekeyEvent(&msg)
+}
+
 func handleUnknownEvent(stream Stream, command string) error {
 	var msg struct{}
 	if err := recvMessage(stream, &msg); err != nil {

tsshd/client.go

@@ -136,11 +136,7 @@ func NewSshUdpClient(opts *UdpClientOptions) (*SshUdpClient, error) {
 		return nil, err
 	}
 
-	mtu := uint16(0)
-	if opts.ProxyClient != nil {
-		mtu = opts.ProxyClient.GetMaxDatagramSize()
-	}
-	udpClient.protoClient, err = newProtoClient(tsshdAddr, opts.ServerInfo, mtu, opts.ConnectTimeout)
+	udpClient.protoClient, err = newProtoClient(udpClient, opts, tsshdAddr)
 	if err != nil {
 		return nil, err
 	}
@@ -426,6 +422,10 @@ func (c *SshUdpClient) GetLastReconnectError() error {
 // GetMaxDatagramSize returns the maximum payload size (in bytes) that
 // can be sent in a single datagram over this SshUdpClient.
 func (c *SshUdpClient) GetMaxDatagramSize() uint16 {
+	return getMaxDatagramSizeFunc(c)
+}
+
+var getMaxDatagramSizeFunc = func(c *SshUdpClient) uint16 {
 	return c.protoClient.getUdpForwarder().conn.GetMaxDatagramSize()
 }
 
@@ -519,6 +519,12 @@ func (c *SshUdpClient) keepAlive(intervalTime time.Duration) {
 	}
 }
 
+func (c *SshUdpClient) isBusStreamInited() bool {
+	c.busMutex.Lock()
+	defer c.busMutex.Unlock()
+	return c.busStream != nil
+}
+
 func (c *SshUdpClient) sendBusCommand(command string) error {
 	c.busMutex.Lock()
 	defer c.busMutex.Unlock()
@@ -556,6 +562,8 @@ func (c *SshUdpClient) handleBusEvent() {
 			c.handleAliveEvent()
 		case "discard":
 			c.handleDiscardEvent()
+		case "rekey":
+			c.handleRekeyEvent()
 		default:
 			if err := handleUnknownEvent(c.busStream, command); err != nil {
 				c.warning("handle bus command [%s] failed: %v. You may need to upgrade tssh.", command, err)
@@ -665,6 +673,19 @@ func (c *SshUdpClient) handleDiscardEvent() {
 	}
 }
 
+func (c *SshUdpClient) handleRekeyEvent() {
+	var msg rekeyMessage
+	if err := recvMessage(c.busStream, &msg); err != nil {
+		c.warning("recv rekey message failed: %v", err)
+		return
+	}
+
+	if err := c.protoClient.handleRekeyEvent(&msg); err != nil {
+		c.warning("rekey failed: %v", err)
+		return
+	}
+}
+
 // SshUdpSession represents a connection to a remote command or shell
 type SshUdpSession struct {
 	id      uint64

tsshd/datagram.go

@@ -35,8 +35,6 @@ import (
 	"sync/atomic"
 )
 
-var globalUdpForwarder *udpForwarder
-
 var udpForwardChannelID atomic.Uint64
 
 // PacketConn represents a connection capable of sending and receiving packet-based data.
@@ -349,7 +347,7 @@ func handleDialUdpEvent(stream Stream) {
 	}
 
 	id := udpForwardChannelID.Add(1)
-	pconn := newPacketConn(stream, id, globalUdpForwarder, globalServerProxy.clientChecker)
+	pconn := newPacketConn(stream, id, globalProtoServer.getUdpForwarder(), globalServerProxy.clientChecker)
 
 	resp := dialUdpResponse{ID: id}
 	if err := sendResponse(stream, &resp); err != nil { // ack ok