go-tunnel/server_config.go

package tunnel

import (
	"crypto/sha1"
	"time"

	"github.com/pkg/errors"
	"github.com/xtaci/kcp-go/v5"
	"github.com/xtaci/smux"
	"golang.org/x/crypto/pbkdf2"
)

type ConfigureConnFunc func(conn *kcp.UDPSession) error
type ConfigureListenerFunc func(listener *kcp.Listener) error

type ServerConfig struct {
	Address               string
	BlockCrypt            kcp.BlockCrypt
	DataShards            int
	ParityShards          int
	Hooks                 *ServerHooks
	ConfigureConn         ConfigureConnFunc
	ConfigureListener     ConfigureListenerFunc
	AuthenticationTimeout time.Duration
	ProxyRequestTimeout   time.Duration
	SmuxConfig            *smux.Config
}

// nolint: go-mnd
func DefaultServerConfig() *ServerConfig {
	unencryptedBlock, err := kcp.NewNoneBlockCrypt(nil)
	if err != nil { // should never happen
		panic(errors.WithStack(err))
	}

	smuxConfig := smux.DefaultConfig()
	smuxConfig.Version = 2
	smuxConfig.KeepAliveInterval = 10 * time.Second
	smuxConfig.MaxReceiveBuffer = 4194304
	smuxConfig.MaxStreamBuffer = 2097152

	return &ServerConfig{
		Address:      ":36543",
		BlockCrypt:   unencryptedBlock,
		DataShards:   3,
		ParityShards: 10,
		Hooks: &ServerHooks{
			onClientConnect:    DefaultOnClientConnect,
			onClientDisconnect: DefaultOnClientDisconnect,
			onClientAuth:       DefaultOnClientAuth,
		},
		ConfigureConn:         DefaultServerConfigureConn,
		ConfigureListener:     DefaultServerConfigureListener,
		AuthenticationTimeout: 30 * time.Second,
		ProxyRequestTimeout:   5 * time.Second,
		SmuxConfig:            smuxConfig,
	}
}

type ServerConfigFunc func(c *ServerConfig)

func WithServerAddress(address string) ServerConfigFunc {
	return func(conf *ServerConfig) {
		conf.Address = address
	}
}

func WithServerBlockCrypt(alg string, pass, salt string, iterations, keyLen int) ServerConfigFunc {
	return func(conf *ServerConfig) {
		key := pbkdf2.Key([]byte(pass), []byte(salt), iterations, keyLen, sha1.New)

		block, err := createBlockCrypt(alg, key)
		if err != nil {
			panic(errors.Wrap(err, "could not create block crypt"))
		}

		conf.BlockCrypt = block
	}
}

func WithServerOnClientAuth(fn OnClientAuthFunc) ServerConfigFunc {
	return func(conf *ServerConfig) {
		conf.Hooks.onClientAuth = fn
	}
}

func WithServerOnClientConnect(fn OnClientConnectFunc) ServerConfigFunc {
	return func(conf *ServerConfig) {
		conf.Hooks.onClientConnect = fn
	}
}

func WithServerOnClientDisconnect(fn OnClientDisconnectFunc) ServerConfigFunc {
	return func(conf *ServerConfig) {
		conf.Hooks.onClientDisconnect = fn
	}
}

func WithServerConfigureConn(fn ConfigureConnFunc) ServerConfigFunc {
	return func(conf *ServerConfig) {
		conf.ConfigureConn = fn
	}
}

func WithServerConfigureListener(fn ConfigureListenerFunc) ServerConfigFunc {
	return func(conf *ServerConfig) {
		conf.ConfigureListener = fn
	}
}

func WithServerSmuxConfig(c *smux.Config) ServerConfigFunc {
	return func(conf *ServerConfig) {
		conf.SmuxConfig = c
	}
}

// nolint: go-mnd
func DefaultServerConfigureConn(conn *kcp.UDPSession) error {
	// Based on kcptun default configuration, mode 'fast3'
	conn.SetStreamMode(true)
	conn.SetWriteDelay(false)
	conn.SetNoDelay(1, 10, 2, 1)
	conn.SetWindowSize(128, 512)
	conn.SetMtu(1400)
	conn.SetACKNoDelay(true)

	if err := conn.SetDSCP(46); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// nolint: go-mnd
func DefaultServerConfigureListener(listener *kcp.Listener) error {
	// Based on kcptun default configuration, mode 'fast3'
	if err := listener.SetReadBuffer(16777217); err != nil {
		return errors.WithStack(err)
	}

	if err := listener.SetWriteBuffer(16777217); err != nil {
		return errors.WithStack(err)
	}

	if err := listener.SetDSCP(46); err != nil {
		return errors.WithStack(err)
	}

	return nil
}