Last active
June 1, 2019 11:18
-
-
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).
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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") | |
}() | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Also does not support
CancelRequest
(support for both it andSSLRequest
seems to be lacking frompgx/pgproto3
).