Wormhole/internal/agent/agent.go

package agent

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/xtaci/smux"
	"github.com/wormhole/wormhole/internal/bridge"
	"github.com/wormhole/wormhole/internal/protocol"
	"github.com/wormhole/wormhole/internal/version"
)

func Run(args []string) {
	fs := flag.NewFlagSet("agent", flag.ExitOnError)
	serverAddr := fs.String("server", "127.0.0.1:8528", "wormhole server bridge address")
	verifyKey := fs.String("key", "", "client verify key")
	useTLS := fs.Bool("tls", false, "use TLS for bridge connection")
	fs.Parse(args)

	if *verifyKey == "" {
		log.Fatal("verify key required: -key")
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

	backoff := time.Second
	for {
		select {
		case <-sig:
			log.Println("[agent] shutting down")
			return
		default:
		}

		shutdown, hadSession := runSession(*serverAddr, *verifyKey, *useTLS, sig)
		if shutdown {
			return
		}
		if hadSession {
			backoff = time.Second
		}

		log.Printf("[agent] reconnecting in %v...", backoff)
		select {
		case <-sig:
			log.Println("[agent] shutting down")
			return
		case <-time.After(backoff):
		}
		if backoff < 30*time.Second {
			backoff *= 2
		}
	}
}

func dialBridge(addr string, useTLS bool) (net.Conn, error) {
	var conn net.Conn
	var err error
	if useTLS {
		conn, err = tls.Dial("tcp", addr, &tls.Config{
			MinVersion:         tls.VersionTLS12,
			InsecureSkipVerify: true,
		})
	} else {
		conn, err = net.Dial("tcp", addr)
	}
	if err != nil {
		return nil, err
	}
	bridge.SetTCPKeepAlive(conn)
	return conn, nil
}

// runSession connects and blocks until disconnect.
// Returns shutdown=true when SIGINT/SIGTERM received; hadSession=true if registration succeeded.
func runSession(serverAddr, verifyKey string, useTLS bool, sig chan os.Signal) (shutdown, hadSession bool) {
	conn, err := dialBridge(serverAddr, useTLS)
	if err != nil {
		log.Printf("[agent] dial server: %v", err)
		return false, false
	}
	defer conn.Close()

	if err := protocol.WriteMessage(conn, protocol.MsgRegister, protocol.RegisterPayload{
		VerifyKey: verifyKey,
		Version:   version.Version,
	}); err != nil {
		log.Printf("[agent] register: %v", err)
		return false, false
	}
	msgType, body, err := protocol.ReadMessage(conn)
	if err != nil {
		log.Printf("[agent] register ack: %v", err)
		return false, false
	}
	if msgType != protocol.MsgRegisterAck {
		log.Printf("[agent] unexpected register response")
		return false, false
	}
	var ack protocol.RegisterAckPayload
	if len(body) > 0 {
		if err := json.Unmarshal(body, &ack); err != nil {
			log.Printf("[agent] parse ack: %v", err)
			return false, false
		}
	}
	if !ack.OK {
		log.Printf("[agent] register failed: %s", ack.Message)
		return false, false
	}
	log.Printf("[agent] registered as client %d", ack.ClientID)
	hadSession = true

	session, err := smux.Server(conn, bridge.SmuxConfig())
	if err != nil {
		log.Printf("[agent] smux: %v", err)
		return false, hadSession
	}
	defer session.Close()

	sessionDead := make(chan string, 1)
	go func() {
		reason := heartbeat(session)
		sessionDead <- reason
	}()

	go func() {
		for {
			stream, err := session.AcceptStream()
			if err != nil {
				select {
				case sessionDead <- "accept stream: " + err.Error():
				default:
				}
				return
			}
			go handleStream(stream)
		}
	}()

	select {
	case <-sig:
		return true, hadSession
	case reason := <-sessionDead:
		log.Printf("[agent] session closed (%s)", reason)
		return false, hadSession
	}
}

const heartbeatMaxFailures = 3

func heartbeat(session *smux.Session) string {
	ticker := time.NewTicker(bridge.HeartbeatInterval)
	defer ticker.Stop()
	failures := 0
	for range ticker.C {
		if err := pingOnce(session); err != nil {
			failures++
			log.Printf("[agent] heartbeat failed (%d/%d): %v", failures, heartbeatMaxFailures, err)
			if failures >= heartbeatMaxFailures {
				return "heartbeat timeout"
			}
			continue
		}
		failures = 0
	}
	return "heartbeat stopped"
}

func pingOnce(session *smux.Session) error {
	stream, err := session.OpenStream()
	if err != nil {
		return err
	}
	defer stream.Close()
	if err := protocol.WriteMessage(stream, protocol.MsgPing, nil); err != nil {
		return err
	}
	msgType, _, err := protocol.ReadMessage(stream)
	if err != nil {
		return err
	}
	if msgType != protocol.MsgPong {
		return fmt.Errorf("unexpected response type %d", msgType)
	}
	return nil
}

func handleStream(stream *smux.Stream) {
	defer stream.Close()
	msgType, body, err := protocol.ReadMessage(stream)
	if err != nil {
		return
	}
	switch msgType {
	case protocol.MsgPing:
		_ = protocol.WriteMessage(stream, protocol.MsgPong, nil)
		return
	case protocol.MsgOpenStream:
		handleDataStream(stream, body)
	}
}

func handleDataStream(stream *smux.Stream, body []byte) {
	meta, err := protocol.ParseLinkMeta(body)
	if err != nil {
		return
	}
	target, err := net.Dial("tcp", meta.Target)
	if err != nil {
		return
	}
	defer target.Close()
	_ = protocol.WriteMessage(stream, protocol.MsgStreamReady, nil)
	go func() {
		_, _ = io.Copy(target, stream)
	}()
	_, _ = io.Copy(stream, target)
}