go-tunnel/server_config.go

package tunnel

import (
	"crypto/sha1"

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

type ConfigureConnFunc func(conn *kcp.UDPSession) error

type ServerConfig struct {
	Address       string
	BlockCrypt    kcp.BlockCrypt
	DataShards    int
	ParityShards  int
	Hooks         *ServerHooks
	ConfigureConn ConfigureConnFunc
}

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

	return &ServerConfig{
		Address:      ":36543",
		BlockCrypt:   unencryptedBlock,
		DataShards:   3,
		ParityShards: 10,
		Hooks: &ServerHooks{
			onClientConnect:    DefaultOnClientConnect,
			onClientDisconnect: DefaultOnClientDisconnect,
			onClientAuth:       DefaultOnClientAuth,
		},
	}
}

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
	}
}