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 ServerConfig struct {
	Address      string
	BlockCrypt   kcp.BlockCrypt
	DataShards   int
	ParityShards int
	Hooks        *ServerHooks
}

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 WithServerAESBlockCrypt(pass, salt string) ServerConfigFunc {
	return func(conf *ServerConfig) {
		key := pbkdf2.Key([]byte(pass), []byte(salt), 1024, 32, sha1.New)

		block, err := kcp.NewAESBlockCrypt(key)
		if err != nil {
			panic(errors.WithStack(err))
		}

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