132 lines
2.8 KiB
Go
Executable File
132 lines
2.8 KiB
Go
Executable File
package src
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"hash/crc32"
|
|
"io"
|
|
"log"
|
|
"net"
|
|
"os"
|
|
"slices"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
type Config struct {
|
|
Listen string `json:"LISTEN"`
|
|
Forwards string `json:"FORWARDS"`
|
|
Hello string `json:"HELLO_B64"`
|
|
forwards []*sync.Pool
|
|
}
|
|
|
|
func NewConfig(ctx context.Context) (Config, error) {
|
|
config := Config{
|
|
Listen: ":10000",
|
|
Forwards: "",
|
|
Hello: base64.StdEncoding.EncodeToString([]byte("*1\r\n$4\r\nping\r\n")),
|
|
}
|
|
b, _ := json.Marshal(config)
|
|
var m map[string]any
|
|
if err := json.Unmarshal(b, &m); err != nil {
|
|
return config, err
|
|
}
|
|
for k := range m {
|
|
if v := os.Getenv(k); v != "" {
|
|
m[k] = v
|
|
}
|
|
}
|
|
b2, _ := json.Marshal(m)
|
|
if err := json.Unmarshal(b2, &config); err != nil {
|
|
return config, err
|
|
}
|
|
|
|
hello, err := base64.StdEncoding.DecodeString(config.Hello)
|
|
if err != nil {
|
|
return config, err
|
|
}
|
|
|
|
forwards := strings.Split(config.Forwards, ",")
|
|
forwards = slices.DeleteFunc(forwards, func(s string) bool { return s == "" })
|
|
if len(forwards) == 0 {
|
|
return config, fmt.Errorf("at least one $FORWARD required")
|
|
}
|
|
config.forwards = make([]*sync.Pool, len(forwards))
|
|
for i := range forwards {
|
|
forward := forwards[i]
|
|
config.forwards[i] = &sync.Pool{
|
|
New: func() any {
|
|
if ctx.Err() != nil {
|
|
return nil
|
|
}
|
|
log.Printf("dialing %q with %q", forward, hello)
|
|
v, err := (&net.Dialer{}).DialContext(ctx, "tcp", forward)
|
|
if err != nil {
|
|
log.Printf("! failed dial %q: %v", forward, err)
|
|
return nil
|
|
}
|
|
if _, err := v.Write([]byte(hello)); err != nil {
|
|
log.Printf("! failed write hello %q: %v", forward, err)
|
|
v.Close()
|
|
return nil
|
|
}
|
|
if raw, _, err := readMessage(bufio.NewReader(v)); err != nil {
|
|
log.Printf("! failed read hello %q: %v", forward, err)
|
|
v.Close()
|
|
return nil
|
|
} else {
|
|
log.Printf("dial reply: %q", raw)
|
|
}
|
|
return v
|
|
},
|
|
}
|
|
}
|
|
|
|
return config, nil
|
|
}
|
|
|
|
func (c Config) Close() {
|
|
for i := range c.forwards {
|
|
if c.forwards[i] != nil {
|
|
i := i
|
|
func() {
|
|
c.forwards[i].New = nil
|
|
for {
|
|
got := c.forwards[i].Get()
|
|
if got == nil {
|
|
break
|
|
}
|
|
got.(net.Conn).Close()
|
|
}
|
|
}()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c Config) WithConn(hashKey string, foo func(net.Conn) error) error {
|
|
hash := int(crc32.ChecksumIEEE([]byte(hashKey)))
|
|
hashIdx := hash % len(c.forwards)
|
|
forward := c.forwards[hashIdx]
|
|
|
|
log.Printf("acquire conn %v", hashIdx)
|
|
forwardCon := forward.Get()
|
|
if forwardCon == nil {
|
|
log.Printf("got a nil conn to %v", hashIdx)
|
|
return io.EOF
|
|
}
|
|
forwardConn := forwardCon.(net.Conn)
|
|
|
|
if err := foo(forwardConn); err != nil {
|
|
log.Printf("errored with conn %v: %v", hashIdx, err)
|
|
forwardConn.Close()
|
|
return err
|
|
}
|
|
|
|
log.Printf("release conn %v", hashIdx)
|
|
forward.Put(forwardConn)
|
|
return nil
|
|
}
|