111 lines
2.2 KiB
Go
111 lines
2.2 KiB
Go
package transport
|
|
|
|
import (
|
|
"agentdb/internal/model"
|
|
"bufio"
|
|
"context"
|
|
"crypto/tls"
|
|
"crypto/x509"
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
func TLSConfig(p model.Proxy) (*tls.Config, error) {
|
|
if !p.TargetTLS {
|
|
return nil, nil
|
|
}
|
|
cfg := &tls.Config{ServerName: p.TargetHost, MinVersion: tls.VersionTLS12}
|
|
if p.TargetCA != "" {
|
|
pool, err := x509.SystemCertPool()
|
|
if err != nil {
|
|
pool = x509.NewCertPool()
|
|
}
|
|
if !pool.AppendCertsFromPEM([]byte(p.TargetCA)) {
|
|
return nil, errors.New("invalid target CA")
|
|
}
|
|
cfg.RootCAs = pool
|
|
}
|
|
return cfg, nil
|
|
}
|
|
func Dial(ctx context.Context, p model.Proxy, dial model.DialFunc) (net.Conn, error) {
|
|
c, err := dial(ctx, "tcp", net.JoinHostPort(p.TargetHost, strconv.Itoa(p.TargetPort)))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
cfg, err := TLSConfig(p)
|
|
if err != nil {
|
|
c.Close()
|
|
return nil, err
|
|
}
|
|
if cfg != nil {
|
|
secure := tls.Client(c, cfg)
|
|
if err = secure.HandshakeContext(ctx); err != nil {
|
|
c.Close()
|
|
return nil, err
|
|
}
|
|
c = secure
|
|
}
|
|
return c, nil
|
|
}
|
|
|
|
type BufferedConn struct {
|
|
net.Conn
|
|
Reader *bufio.Reader
|
|
}
|
|
|
|
func (c *BufferedConn) Read(p []byte) (int, error) { return c.Reader.Read(p) }
|
|
func Accept(ctx context.Context, raw net.Conn, cfg *tls.Config) (net.Conn, error) {
|
|
raw.SetDeadline(time.Now().Add(15 * time.Second))
|
|
reader := bufio.NewReader(raw)
|
|
first, err := reader.Peek(1)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var c net.Conn = &BufferedConn{Conn: raw, Reader: reader}
|
|
if first[0] == 22 {
|
|
if cfg == nil {
|
|
return nil, errors.New("gateway TLS certificate missing")
|
|
}
|
|
config := cfg.Clone()
|
|
config.MinVersion = tls.VersionTLS12
|
|
secure := tls.Server(c, config)
|
|
if err = secure.HandshakeContext(ctx); err != nil {
|
|
return nil, err
|
|
}
|
|
c = secure
|
|
}
|
|
return c, nil
|
|
}
|
|
func Write(w io.Writer, b []byte) error {
|
|
for len(b) > 0 {
|
|
n, err := w.Write(b)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if n == 0 {
|
|
return io.ErrShortWrite
|
|
}
|
|
b = b[n:]
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type LockedWriter struct {
|
|
mu sync.Mutex
|
|
Writer io.Writer
|
|
}
|
|
|
|
func (w *LockedWriter) Write(b []byte) (int, error) {
|
|
w.mu.Lock()
|
|
defer w.mu.Unlock()
|
|
err := Write(w.Writer, b)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return len(b), nil
|
|
}
|