104 lines
3.3 KiB
Go
104 lines
3.3 KiB
Go
package mongoproxy
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
"linkport/internal/model"
|
|
"linkport/internal/outbound"
|
|
"time"
|
|
)
|
|
|
|
var _ model.SessionPreparer = (*Handler)(nil)
|
|
|
|
type preparedConnection struct {
|
|
conn *upstream
|
|
target model.Proxy
|
|
}
|
|
|
|
func (v *preparedConnection) matches(p model.Proxy) bool {
|
|
t := v.target
|
|
return t.TargetHost == p.TargetHost && t.TargetPort == p.TargetPort &&
|
|
t.TargetID == p.TargetID && t.TargetAccountID == p.TargetAccountID &&
|
|
t.Outbound == p.Outbound && t.TargetTLS == p.TargetTLS && t.TargetCA == p.TargetCA &&
|
|
t.TargetCompression == p.TargetCompression &&
|
|
t.ConnectionOptions.Charset == p.ConnectionOptions.Charset &&
|
|
t.ConnectionOptions.ApplicationName == p.ConnectionOptions.ApplicationName &&
|
|
t.ConnectionOptions.QueryTimeoutSeconds == p.ConnectionOptions.QueryTimeoutSeconds &&
|
|
t.TargetUser == p.TargetUser && t.TargetPassword == p.TargetPassword &&
|
|
t.Database == p.Database && t.AuthSource == p.AuthSource && t.AccessMode == p.AccessMode
|
|
}
|
|
|
|
// PrepareSession authenticates without a downstream client. Only unused sessions
|
|
// may be handed off; the owner must never return a checked-out session to the pool.
|
|
func (h *Handler) PrepareSession(ctx context.Context, p model.Proxy, dial model.DialFunc) (*model.PreparedSession, error) {
|
|
if p.ClusterEnabled() {
|
|
return nil, errors.New("prepared sessions require a single target")
|
|
}
|
|
if err := p.CheckTargetCredential(); err != nil {
|
|
return nil, err
|
|
}
|
|
p.PreparedTarget = nil
|
|
database := p.AuthSource
|
|
if database == "" {
|
|
database = "admin"
|
|
}
|
|
ctx, cancel := context.WithTimeout(ctx, p.ConnectionOptions.ConnectTimeout(15*time.Second))
|
|
defer cancel()
|
|
c, err := connect(ctx, p, dial)
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return nil, ctx.Err()
|
|
}
|
|
return nil, preparedSetupError(err)
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
_ = c.Close()
|
|
return nil, err
|
|
}
|
|
return &model.PreparedSession{
|
|
Value: &preparedConnection{conn: c, target: p},
|
|
Close: func() { _ = c.Close() },
|
|
Ping: func(ctx context.Context) error {
|
|
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
|
defer cancel()
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
// Wait for an already-running cancellation callback before returning ownership.
|
|
finished := make(chan struct{})
|
|
stop := context.AfterFunc(ctx, func() { _ = c.Close(); close(finished) })
|
|
reply, err := exchange(c, bson.D{{Key: "ping", Value: int32(1)}, {Key: "$db", Value: database}})
|
|
if err == nil && !ok(reply) {
|
|
err = errors.New("MongoDB heartbeat rejected")
|
|
}
|
|
if !stop() {
|
|
<-finished
|
|
}
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
if err != nil {
|
|
return errors.New("prepared target heartbeat failed")
|
|
}
|
|
return nil
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
// Preserve only bounded typed rejection categories; discard all server text.
|
|
func preparedSetupError(err error) error {
|
|
safe := errors.New("prepared target session setup failed")
|
|
for cause := err; cause != nil; cause = errors.Unwrap(cause) {
|
|
switch outbound.FailureKind(cause) {
|
|
case "authentication":
|
|
return model.AtConnectionStage("authentication_rejected", safe)
|
|
case "permission":
|
|
return model.AtConnectionStage("permission_rejected", safe)
|
|
case "identity":
|
|
return model.AtConnectionStage("ssh_host_identity", safe)
|
|
}
|
|
}
|
|
return safe
|
|
}
|