1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
package main
import (
"fmt"
"net"
"github.com/jackc/pgproto3/v2"
)
type PgFortuneBackend struct {
backend *pgproto3.Backend
conn net.Conn
responder func() ([]byte, error)
}
func NewPgFortuneBackend(conn net.Conn, responder func() ([]byte, error)) *PgFortuneBackend {
backend := pgproto3.NewBackend(pgproto3.NewChunkReader(conn), conn)
connHandler := &PgFortuneBackend{
backend: backend,
conn: conn,
responder: responder,
}
return connHandler
}
func (p *PgFortuneBackend) Run() error {
defer p.Close()
err := p.handleStartup()
if err != nil {
return err
}
for {
msg, err := p.backend.Receive()
if err != nil {
return fmt.Errorf("error receiving message: %w", err)
}
switch msg.(type) {
case *pgproto3.Query:
response, err := p.responder()
if err != nil {
return fmt.Errorf("error generating query response: %w", err)
}
buf := mustEncode((&pgproto3.RowDescription{Fields: []pgproto3.FieldDescription{
{
Name: []byte("fortune"),
TableOID: 0,
TableAttributeNumber: 0,
DataTypeOID: 25,
DataTypeSize: -1,
TypeModifier: -1,
Format: 0,
},
}}).Encode(nil))
buf = mustEncode((&pgproto3.DataRow{Values: [][]byte{response}}).Encode(buf))
buf = mustEncode((&pgproto3.CommandComplete{CommandTag: []byte("SELECT 1")}).Encode(buf))
buf = mustEncode((&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(buf))
_, err = p.conn.Write(buf)
if err != nil {
return fmt.Errorf("error writing query response: %w", err)
}
case *pgproto3.Terminate:
return nil
default:
return fmt.Errorf("received message other than Query from client: %#v", msg)
}
}
}
func (p *PgFortuneBackend) handleStartup() error {
startupMessage, err := p.backend.ReceiveStartupMessage()
if err != nil {
return fmt.Errorf("error receiving startup message: %w", err)
}
switch startupMessage.(type) {
case *pgproto3.StartupMessage:
buf := mustEncode((&pgproto3.AuthenticationOk{}).Encode(nil))
buf = mustEncode((&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(buf))
_, err = p.conn.Write(buf)
if err != nil {
return fmt.Errorf("error sending ready for query: %w", err)
}
case *pgproto3.SSLRequest:
_, err = p.conn.Write([]byte("N"))
if err != nil {
return fmt.Errorf("error sending deny SSL request: %w", err)
}
return p.handleStartup()
default:
return fmt.Errorf("unknown startup message: %#v", startupMessage)
}
return nil
}
func (p *PgFortuneBackend) Close() error {
return p.conn.Close()
}
func mustEncode(buf []byte, err error) []byte {
if err != nil {
panic(err)
}
return buf
}
|
