hysteria/extras/network/client.go

package network

import (
	"context"
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/apernet/quic-go"
)

type NetworkClient struct {
	config     UpstreamConfig
	tun        *TUNInterface
	assignedIP net.IP
	gateway    net.IP

	stream *quic.Stream
	conn   *quic.Conn
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

func NewNetworkClient(config UpstreamConfig, tunName string, mtu int) (*NetworkClient, error) {
	if tunName == "" {
		tunName = defaultTUNName + "_c"
	}
	if mtu <= 0 {
		mtu = defaultTUNMTU
	}
	ctx, cancel := context.WithCancel(context.Background())
	return &NetworkClient{
		config: config,
		ctx:    ctx,
		cancel: cancel,
	}, nil
}

func (c *NetworkClient) Connect() error {
	tlsCfg := c.config.TLS
	if tlsCfg == nil {
		var err error
		tlsCfg, err = GenerateNetworkTLSConfig()
		if err != nil {
			return fmt.Errorf("generate TLS: %w", err)
		}
		tlsCfg.InsecureSkipVerify = true
	}
	conn, err := quic.DialAddr(c.ctx, c.config.Server, tlsCfg, nil)
	if err != nil {
		return fmt.Errorf("dial: %w", err)
	}
	c.conn = conn

	stream, err := conn.OpenStream()
	if err != nil {
		conn.CloseWithError(1, "stream error")
		return fmt.Errorf("open stream: %w", err)
	}
	c.stream = stream

	if err := SendAuth(stream, c.config.Token); err != nil {
		return fmt.Errorf("send auth: %w", err)
	}

	f, err := ReceiveFrame(stream)
	if err != nil {
		return fmt.Errorf("recv auth: %w", err)
	}
	switch f.Type {
	case FrameAuthOK:
		if len(f.Payload) < 4 {
			return fmt.Errorf("invalid IP in auth response")
		}
		c.assignedIP = net.IP(f.Payload[:4])
		c.gateway = make(net.IP, 4)
		copy(c.gateway, c.assignedIP)
		c.gateway[3] = 1
		log.Printf("network: assigned IP %s, gateway %s", c.assignedIP, c.gateway)
	case FrameAuthErr:
		return fmt.Errorf("auth rejected: %s", string(f.Payload))
	default:
		return fmt.Errorf("unexpected auth response: %d", f.Type)
	}

	c.tun, err = OpenTUN(defaultTUNName+"_c", c.assignedIP.String()+"/32", defaultTUNMTU)
	if err != nil {
		return fmt.Errorf("open TUN: %w", err)
	}
	c.tun.Route(c.assignedIP.String() + "/32")

	go func() {
		ticker := time.NewTicker(defaultKeepalive)
		defer ticker.Stop()
		for {
			select {
			case <-c.ctx.Done():
				return
			case <-ticker.C:
				SendKeepalive(stream)
			}
		}
	}()

	c.wg.Add(2)
	go c.tunToStream()
	go c.streamToTun()

	return nil
}

func (c *NetworkClient) tunToStream() {
	defer c.wg.Done()
	for {
		pkt, err := c.tun.Read()
		if err != nil {
			if c.ctx.Err() != nil {
				return
			}
			log.Printf("TUN read: %v", err)
			return
		}
		if err := SendData(c.stream, pkt); err != nil {
			log.Printf("send: %v", err)
			return
		}
	}
}

func (c *NetworkClient) streamToTun() {
	defer c.wg.Done()
	for {
		f, err := ReceiveFrame(c.stream)
		if err != nil {
			if c.ctx.Err() != nil {
				return
			}
			log.Printf("recv: %v", err)
			return
		}
		switch f.Type {
		case FrameData:
			if err := c.tun.Write(f.Payload); err != nil {
				log.Printf("TUN write: %v", err)
				return
			}
		case FrameKeepalive:
		default:
		}
	}
}

func (c *NetworkClient) AssignedIP() net.IP {
	return c.assignedIP
}

func (c *NetworkClient) Done() <-chan struct{} {
	return c.ctx.Done()
}

func (c *NetworkClient) Close() error {
	c.cancel()
	if c.stream != nil {
		c.stream.Close()
	}
	if c.conn != nil {
		c.conn.CloseWithError(0, "bye")
	}
	if c.tun != nil {
		c.tun.Close()
	}
	c.wg.Wait()
	return nil
}