Skip to content

Instantly share code, notes, and snippets.

@asterite3
Last active June 1, 2019 11:18
Show Gist options
  • Save asterite3/4ba7824b73e4103b3e5a2c2ef51d6f60 to your computer and use it in GitHub Desktop.
Save asterite3/4ba7824b73e4103b3e5a2c2ef51d6f60 to your computer and use it in GitHub Desktop.
PostgreSQL wire protocol proxy/MitM written in Go based on pgx (https://github.com/jackc/pgx). Does not support SSLRequest (connect with sslmode=disable).
package main
import (
"encoding/json"
"fmt"
"io"
"log"
"net"
"sync"
"github.com/jackc/pgx/pgproto3"
)
const (
bindAddr = "127.0.0.1:5432"
upstreamAddr = "172.17.0.2:5432"
verbose = true
)
func toJSONOrFmt(v interface{}) string {
if _, ok := v.(*pgproto3.Bind); ok {
return fmt.Sprintf("%+v", v)
}
b, err := json.Marshal(v)
if err != nil {
return fmt.Sprintf("%+v", v)
}
return string(b)
}
func main() {
ln, err := net.Listen("tcp", bindAddr)
if err != nil {
panic(err)
}
for {
conn, err := ln.Accept()
if err != nil {
panic(err)
}
go func() {
defer conn.Close()
backend, err := pgproto3.NewBackend(conn, conn)
if err != nil {
panic(err)
}
serverConn, err := net.Dial("tcp", upstreamAddr)
if err != nil {
panic(err)
}
defer serverConn.Close()
frontend, err := pgproto3.NewFrontend(serverConn, serverConn)
if err != nil {
panic(err)
}
exiting := false
log.Printf("New proxy connection started")
wg := &sync.WaitGroup{}
wg.Add(2)
go func() {
// server -> client
defer wg.Done()
for {
msg, err := frontend.Receive()
switch {
case err == nil:
case err == io.EOF && exiting:
return
case err == io.EOF && !exiting:
log.Printf("Connection unexpectedly closed by server")
exiting = true
conn.(*net.TCPConn).CloseWrite()
return
case err != nil && err != io.EOF:
panic(err)
}
if verbose {
log.Printf("Msg from server: %s", toJSONOrFmt(msg))
}
err = backend.Send(msg)
if err != nil {
panic(err)
}
}
}()
go func() {
// client -> server
defer wg.Done()
start, err := backend.ReceiveStartupMessage()
if verbose {
log.Printf("Startup msg from client: %s", toJSONOrFmt(start))
}
if err != nil {
panic(err)
}
err = frontend.Send(start)
if err != nil {
panic(err)
}
for {
msg, err := backend.Receive()
switch {
case err == nil:
case err == io.EOF && exiting:
return
case err == io.EOF && !exiting:
log.Printf("Connection unexpectedly closed by client")
exiting = true
serverConn.(*net.TCPConn).CloseWrite()
return
case err != nil && err != io.EOF:
panic(err)
}
if _, ok := msg.(*pgproto3.Terminate); ok {
exiting = true
}
if verbose {
log.Printf("Msg from client: %s", toJSONOrFmt(msg))
}
err = frontend.Send(msg)
if err != nil {
panic(err)
}
}
}()
wg.Wait()
log.Printf("Proxy connection done")
}()
}
}
@asterite3
Copy link
Author

asterite3 commented Jun 1, 2019

Also does not support CancelRequest (support for both it and SSLRequest seems to be lacking from pgx/pgproto3).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment