281 lines
8.1 KiB
Go
281 lines
8.1 KiB
Go
package pgproxy
|
|
|
|
import (
|
|
"context"
|
|
"crypto/subtle"
|
|
"crypto/tls"
|
|
"errors"
|
|
"fmt"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
"github.com/jackc/pgx/v5/pgproto3"
|
|
"linkport/internal/cluster"
|
|
"linkport/internal/model"
|
|
"linkport/internal/outbound"
|
|
"linkport/internal/protocol/transport"
|
|
"net"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type Handler struct{}
|
|
|
|
func New() *Handler { return &Handler{} }
|
|
func open(ctx context.Context, p model.Proxy, dial model.DialFunc, database string) (*pgconn.PgConn, error) {
|
|
if err := p.CheckTargetCredential(); err != nil {
|
|
return nil, err
|
|
}
|
|
initial := p.AuthSource
|
|
if initial == "" {
|
|
initial = "postgres"
|
|
}
|
|
if !p.ClusterEnabled() && (database == "" || database == initial) {
|
|
if session, ok := p.BorrowPrepared(); ok {
|
|
if v, ok := session.Value.(*preparedConnection); ok && v.matches(p) {
|
|
if err := session.Ping(ctx); err == nil {
|
|
return v.conn, nil
|
|
}
|
|
}
|
|
session.Close()
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
}
|
|
var resolveErr error
|
|
p, _, resolveErr = cluster.Resolve(ctx, p, dial)
|
|
if resolveErr != nil {
|
|
return nil, resolveErr
|
|
}
|
|
if err := p.CheckTargetCredential(); err != nil {
|
|
return nil, err
|
|
}
|
|
cfg, err := pgconn.ParseConfig("postgres://localhost/postgres?sslmode=disable")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
cfg.Host, cfg.Port, cfg.User, cfg.Password = p.TargetHost, uint16(p.TargetPort), p.TargetUser, p.TargetPassword
|
|
cfg.Database = database
|
|
if cfg.Database == "" {
|
|
cfg.Database = p.AuthSource
|
|
}
|
|
if cfg.Database == "" {
|
|
cfg.Database = "postgres"
|
|
}
|
|
cfg.TLSConfig, err = transport.TLSConfig(p)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
cfg.Fallbacks = nil
|
|
cfg.ConnectTimeout = p.ConnectionOptions.ConnectTimeout(15 * time.Second)
|
|
cfg.DialFunc = pgconn.DialFunc(dial)
|
|
// Keep DNS resolution on the configured proxy/jump host.
|
|
cfg.LookupFunc = func(_ context.Context, host string) ([]string, error) { return []string{host}, nil }
|
|
cfg.RuntimeParams = map[string]string{"application_name": "linkport", "client_encoding": "UTF8"}
|
|
if p.ConnectionOptions.ApplicationName != "" {
|
|
cfg.RuntimeParams["application_name"] = p.ConnectionOptions.ApplicationName
|
|
}
|
|
if p.ConnectionOptions.QueryTimeoutSeconds > 0 {
|
|
cfg.RuntimeParams["statement_timeout"] = strconv.Itoa(p.ConnectionOptions.QueryTimeoutSeconds * 1000)
|
|
}
|
|
return pgconn.ConnectConfig(ctx, cfg)
|
|
}
|
|
func (h *Handler) Test(ctx context.Context, p model.Proxy, dial model.DialFunc) (model.TestResult, error) {
|
|
done := outbound.Begin(ctx, "protocol_session")
|
|
c, err := open(ctx, p, dial, "")
|
|
done(err)
|
|
if err != nil {
|
|
return model.TestResult{}, err
|
|
}
|
|
defer c.Close(context.Background())
|
|
done = outbound.Begin(ctx, "target_query")
|
|
results, err := c.Exec(ctx, "SELECT 1").ReadAll()
|
|
if err != nil || len(results) != 1 {
|
|
if err == nil {
|
|
err = errors.New("unexpected query result")
|
|
}
|
|
done(err)
|
|
return model.TestResult{}, errors.New("PostgreSQL query test failed")
|
|
}
|
|
done(nil)
|
|
return model.TestResult{Version: c.ParameterStatus("server_version"), TLS: p.TargetTLS, Compression: "none"}, nil
|
|
}
|
|
func (h *Handler) Serve(ctx context.Context, raw net.Conn, p model.Proxy, dial model.DialFunc, tc *tls.Config, event model.EventFunc) error {
|
|
defer raw.Close()
|
|
stop := context.AfterFunc(ctx, func() { raw.Close() })
|
|
defer stop()
|
|
raw.SetDeadline(time.Now().Add(20 * time.Second))
|
|
var c net.Conn = raw
|
|
backend := pgproto3.NewBackend(c, c)
|
|
backend.SetMaxBodyLen(32 << 20)
|
|
var startup *pgproto3.StartupMessage
|
|
for tries := 0; tries < 3; tries++ {
|
|
msg, err := backend.ReceiveStartupMessage()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
switch m := msg.(type) {
|
|
case *pgproto3.SSLRequest:
|
|
if tc == nil {
|
|
if err = transport.Write(c, []byte("N")); err != nil {
|
|
return err
|
|
}
|
|
continue
|
|
}
|
|
if _, already := c.(*tls.Conn); already {
|
|
return errors.New("duplicate PostgreSQL SSL request")
|
|
}
|
|
if err = transport.Write(c, []byte("S")); err != nil {
|
|
return err
|
|
}
|
|
cfg := tc.Clone()
|
|
cfg.MinVersion = tls.VersionTLS12
|
|
secure := tls.Server(c, cfg)
|
|
if err = secure.HandshakeContext(ctx); err != nil {
|
|
return err
|
|
}
|
|
c = secure
|
|
backend = pgproto3.NewBackend(c, c)
|
|
backend.SetMaxBodyLen(32 << 20)
|
|
case *pgproto3.GSSEncRequest:
|
|
if err = transport.Write(c, []byte("N")); err != nil {
|
|
return err
|
|
}
|
|
case *pgproto3.StartupMessage:
|
|
startup = m
|
|
default:
|
|
return errors.New("unsupported PostgreSQL startup request")
|
|
}
|
|
if startup != nil {
|
|
break
|
|
}
|
|
}
|
|
if startup == nil {
|
|
return errors.New("missing PostgreSQL startup")
|
|
}
|
|
fail := func(message string) error {
|
|
backend.Send(&pgproto3.ErrorResponse{Severity: "FATAL", Code: "28000", Message: message})
|
|
backend.Flush()
|
|
return errors.New(message)
|
|
}
|
|
username := startup.Parameters["user"]
|
|
account, exists := transport.Account(p, username, raw.RemoteAddr())
|
|
if !exists {
|
|
transport.AuthenticationEvent(event, username, false)
|
|
return fail("gateway authentication failed")
|
|
}
|
|
p = account.Apply(p)
|
|
dial = p.DialForTarget(dial)
|
|
backend.SetAuthType(pgproto3.AuthTypeCleartextPassword)
|
|
backend.Send(&pgproto3.AuthenticationCleartextPassword{})
|
|
if err := backend.Flush(); err != nil {
|
|
return err
|
|
}
|
|
msg, err := backend.Receive()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
password, ok := msg.(*pgproto3.PasswordMessage)
|
|
if !ok || subtle.ConstantTimeCompare([]byte(password.Password), []byte(p.GatewayPassword)) != 1 {
|
|
transport.AuthenticationEvent(event, username, false)
|
|
return fail("gateway authentication failed")
|
|
}
|
|
transport.AuthenticationEvent(event, username, true)
|
|
upstream, err := open(ctx, p, dial, startup.Parameters["database"])
|
|
if err != nil {
|
|
return fail("target connection or authentication failed")
|
|
}
|
|
if err = upstream.SyncConn(ctx); err != nil {
|
|
upstream.Close(ctx)
|
|
return err
|
|
}
|
|
connection, err := upstream.Hijack()
|
|
if err != nil {
|
|
upstream.Close(ctx)
|
|
return err
|
|
}
|
|
defer connection.Conn.Close()
|
|
stopUp := context.AfterFunc(ctx, func() { connection.Conn.Close() })
|
|
defer stopUp()
|
|
backend.Send(&pgproto3.AuthenticationOk{})
|
|
for name, value := range connection.ParameterStatuses {
|
|
backend.Send(&pgproto3.ParameterStatus{Name: name, Value: value})
|
|
}
|
|
backend.Send(&pgproto3.ReadyForQuery{TxStatus: connection.TxStatus})
|
|
if err = backend.Flush(); err != nil {
|
|
return err
|
|
}
|
|
raw.SetDeadline(time.Time{})
|
|
connection.Conn.SetDeadline(time.Time{})
|
|
done := make(chan error, 1)
|
|
var timing model.ExecutionQueue
|
|
go func() {
|
|
for {
|
|
message, e := connection.Frontend.Receive()
|
|
if e != nil {
|
|
done <- e
|
|
raw.Close()
|
|
return
|
|
}
|
|
backend.Send(message)
|
|
if e = backend.Flush(); e != nil {
|
|
done <- e
|
|
raw.Close()
|
|
return
|
|
}
|
|
if _, ok := message.(*pgproto3.ReadyForQuery); ok {
|
|
timing.Finish(ctx, "pgsql")
|
|
}
|
|
}
|
|
}()
|
|
var extendedStarted time.Time
|
|
for {
|
|
message, err := backend.Receive()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
kind, text := "", ""
|
|
switch m := message.(type) {
|
|
case *pgproto3.Query:
|
|
kind, text = "query", m.String
|
|
case *pgproto3.Parse:
|
|
kind, text = "prepare", m.Query
|
|
case *pgproto3.Bind:
|
|
kind, text = "bind", fmt.Sprintf("statement=%s portal=%s parameters=%d", m.PreparedStatement, m.DestinationPortal, len(m.Parameters))
|
|
case *pgproto3.Execute:
|
|
kind, text = "execute", fmt.Sprintf("portal=%s maxRows=%d", m.Portal, m.MaxRows)
|
|
case *pgproto3.Terminate:
|
|
kind, text = "quit", "terminate"
|
|
case *pgproto3.CopyData:
|
|
kind, text = "copy_data", fmt.Sprintf("bytes=%d", len(m.Data))
|
|
default:
|
|
kind = strings.TrimPrefix(fmt.Sprintf("%T", message), "*pgproto3.")
|
|
text = kind
|
|
}
|
|
if err = model.AuditCommand(ctx, kind, text, "allowed"); err != nil {
|
|
return err
|
|
}
|
|
switch message.(type) {
|
|
case *pgproto3.Query:
|
|
timing.Start("query batch")
|
|
case *pgproto3.Parse, *pgproto3.Bind, *pgproto3.Execute:
|
|
if extendedStarted.IsZero() {
|
|
extendedStarted = time.Now()
|
|
}
|
|
case *pgproto3.Sync:
|
|
if extendedStarted.IsZero() {
|
|
extendedStarted = time.Now()
|
|
}
|
|
timing.StartAt("extended query cycle", extendedStarted)
|
|
extendedStarted = time.Time{}
|
|
}
|
|
connection.Frontend.Send(message)
|
|
if err = connection.Frontend.Flush(); err != nil {
|
|
return err
|
|
}
|
|
if _, ok := message.(*pgproto3.Terminate); ok {
|
|
return nil
|
|
}
|
|
}
|
|
}
|