From 66886d3ef9de6ebb2fca665789cc92c3bee8a01b Mon Sep 17 00:00:00 2001
From: Garet Halliday <me@garet.holiday>
Date: Thu, 24 Aug 2023 14:54:46 -0500
Subject: [PATCH] sizes
---
lib/bouncer/backends/v0/accept.go | 13 ++++++++----
lib/bouncer/backends/v0/cancel.go | 2 +-
lib/zap/packet.go | 20 +++++++++++++++++--
.../packets/v3.0/authenticationcleartext.go | 2 +-
lib/zap/packets/v3.0/authenticationmd5.go | 2 +-
lib/zap/packets/v3.0/authenticationok.go | 2 +-
.../packets/v3.0/authenticationresponse.go | 2 +-
lib/zap/packets/v3.0/authenticationsasl.go | 8 +++++++-
.../v3.0/authenticationsaslcontinue.go | 2 +-
.../packets/v3.0/authenticationsaslfinal.go | 2 +-
lib/zap/packets/v3.0/backendkeydata.go | 2 +-
lib/zap/packets/v3.0/bind.go | 14 ++++++++++++-
lib/zap/packets/v3.0/close.go | 2 +-
lib/zap/packets/v3.0/copyfail.go | 2 +-
lib/zap/packets/v3.0/describe.go | 2 +-
lib/zap/packets/v3.0/errorresponse.go | 10 +++++++++-
lib/zap/packets/v3.0/execute.go | 2 +-
.../packets/v3.0/negotiateprotocolversion.go | 7 ++++++-
lib/zap/packets/v3.0/parameterstatus.go | 2 +-
lib/zap/packets/v3.0/parse.go | 2 +-
lib/zap/packets/v3.0/passwordmessage.go | 2 +-
lib/zap/packets/v3.0/query.go | 2 +-
lib/zap/packets/v3.0/readyforquery.go | 2 +-
lib/zap/packets/v3.0/saslinitialresponse.go | 2 +-
24 files changed, 80 insertions(+), 28 deletions(-)
diff --git a/lib/bouncer/backends/v0/accept.go b/lib/bouncer/backends/v0/accept.go
index 7e13419e..2ea282d9 100644
--- a/lib/bouncer/backends/v0/accept.go
+++ b/lib/bouncer/backends/v0/accept.go
@@ -231,7 +231,7 @@ func startup1(conn *bouncer.Conn) (done bool, err error) {
}
func enableSSL(server zap.ReadWriter, config *tls.Config) (bool, error) {
- packet := zap.NewPacket(0)
+ packet := zap.NewPacket(0, 4)
packet = packet.AppendUint16(1234)
packet = packet.AppendUint16(5679)
if err := server.WritePacket(packet); err != nil {
@@ -280,8 +280,13 @@ func Accept(server zap.ReadWriter, options AcceptOptions) (bouncer.Conn, error)
}
}
- // we can re-use the memory for this pkt most of the way down because we don't pass this anywhere
- packet := zap.NewPacket(0)
+ size := 4 + len("user") + 1 + len(username) + 1 + len("database") + 1 + len(options.Database) + 1
+ for key, value := range options.StartupParameters {
+ size += len(key.String()) + len(value) + 2
+ }
+ size += 1
+
+ packet := zap.NewPacket(0, size)
packet = packet.AppendUint16(3)
packet = packet.AppendUint16(0)
packet = packet.AppendString("user")
@@ -292,7 +297,7 @@ func Accept(server zap.ReadWriter, options AcceptOptions) (bouncer.Conn, error)
packet = packet.AppendString(key.String())
packet = packet.AppendString(value)
}
- packet = packet.AppendString("")
+ packet = packet.AppendUint8(0)
err := server.WritePacket(packet)
if err != nil {
diff --git a/lib/bouncer/backends/v0/cancel.go b/lib/bouncer/backends/v0/cancel.go
index 06f73845..e5659b2b 100644
--- a/lib/bouncer/backends/v0/cancel.go
+++ b/lib/bouncer/backends/v0/cancel.go
@@ -3,7 +3,7 @@ package backends
import "pggat2/lib/zap"
func Cancel(server zap.ReadWriter, key [8]byte) error {
- packet := zap.NewPacket(0)
+ packet := zap.NewPacket(0, 12)
packet = packet.AppendUint16(1234)
packet = packet.AppendUint16(5678)
packet = packet.AppendBytes(key[:])
diff --git a/lib/zap/packet.go b/lib/zap/packet.go
index 88e275dc..73d12d39 100644
--- a/lib/zap/packet.go
+++ b/lib/zap/packet.go
@@ -7,8 +7,14 @@ import (
type Packet []byte
-func NewPacket(typ Type) Packet {
- return []byte{byte(typ), 0, 0, 0, 4}
+func NewPacket(typ Type, size ...int) Packet {
+ if len(size) > 0 {
+ packet := make([]byte, 5, 5+size[0])
+ packet[0] = byte(typ)
+ return packet
+ } else {
+ return []byte{byte(typ), 0, 0, 0, 0}
+ }
}
func (T Packet) Payload() PacketFragment {
@@ -39,6 +45,16 @@ func (T Packet) Grow(n int) Packet {
return p
}
+func (T Packet) Reserve(n int) Packet {
+ if len(T)+n > cap(T) {
+ p := make([]byte, len(T), len(T)+n)
+ copy(p, T)
+ return p
+ }
+
+ return T
+}
+
func (T Packet) Type() Type {
return Type(T[0])
}
diff --git a/lib/zap/packets/v3.0/authenticationcleartext.go b/lib/zap/packets/v3.0/authenticationcleartext.go
index 9985942a..027e24e0 100644
--- a/lib/zap/packets/v3.0/authenticationcleartext.go
+++ b/lib/zap/packets/v3.0/authenticationcleartext.go
@@ -17,7 +17,7 @@ func (T *AuthenticationCleartext) ReadFromPacket(packet zap.Packet) bool {
}
func (T *AuthenticationCleartext) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthentication)
+ packet := zap.NewPacket(TypeAuthentication, 4)
packet = packet.AppendUint32(3)
return packet
}
diff --git a/lib/zap/packets/v3.0/authenticationmd5.go b/lib/zap/packets/v3.0/authenticationmd5.go
index fd480670..8f7ad4e9 100644
--- a/lib/zap/packets/v3.0/authenticationmd5.go
+++ b/lib/zap/packets/v3.0/authenticationmd5.go
@@ -20,7 +20,7 @@ func (T *AuthenticationMD5) ReadFromPacket(packet zap.Packet) bool {
}
func (T *AuthenticationMD5) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthentication)
+ packet := zap.NewPacket(TypeAuthentication, 8)
packet = packet.AppendUint32(5)
packet = packet.AppendBytes(T.Salt[:])
return packet
diff --git a/lib/zap/packets/v3.0/authenticationok.go b/lib/zap/packets/v3.0/authenticationok.go
index e73fc9fa..73e4ac45 100644
--- a/lib/zap/packets/v3.0/authenticationok.go
+++ b/lib/zap/packets/v3.0/authenticationok.go
@@ -17,7 +17,7 @@ func (T *AuthenticationOk) ReadFromPacket(packet zap.Packet) bool {
}
func (T *AuthenticationOk) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthentication)
+ packet := zap.NewPacket(TypeAuthentication, 4)
packet = packet.AppendUint32(0)
return packet
}
diff --git a/lib/zap/packets/v3.0/authenticationresponse.go b/lib/zap/packets/v3.0/authenticationresponse.go
index 260be177..98ab5164 100644
--- a/lib/zap/packets/v3.0/authenticationresponse.go
+++ b/lib/zap/packets/v3.0/authenticationresponse.go
@@ -17,6 +17,6 @@ func (T *AuthenticationResponse) ReadFromPacket(packet zap.Packet) bool {
}
func (T *AuthenticationResponse) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthenticationResponse)
+ packet := zap.NewPacket(TypeAuthenticationResponse, len(*T))
return packet.AppendBytes(*T)
}
diff --git a/lib/zap/packets/v3.0/authenticationsasl.go b/lib/zap/packets/v3.0/authenticationsasl.go
index 983306ff..3f60b9b7 100644
--- a/lib/zap/packets/v3.0/authenticationsasl.go
+++ b/lib/zap/packets/v3.0/authenticationsasl.go
@@ -28,7 +28,13 @@ func (T *AuthenticationSASL) ReadFromPacket(packet zap.Packet) bool {
}
func (T *AuthenticationSASL) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthentication)
+ size := 5
+ for _, mechanism := range T.Mechanisms {
+ size += len(mechanism) + 1
+ }
+
+ packet := zap.NewPacket(TypeAuthentication, size)
+
packet = packet.AppendInt32(10)
for _, mechanism := range T.Mechanisms {
packet = packet.AppendString(mechanism)
diff --git a/lib/zap/packets/v3.0/authenticationsaslcontinue.go b/lib/zap/packets/v3.0/authenticationsaslcontinue.go
index 4e5b1d79..99f616e1 100644
--- a/lib/zap/packets/v3.0/authenticationsaslcontinue.go
+++ b/lib/zap/packets/v3.0/authenticationsaslcontinue.go
@@ -22,7 +22,7 @@ func (T *AuthenticationSASLContinue) ReadFromPacket(packet zap.Packet) bool {
}
func (T *AuthenticationSASLContinue) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthentication)
+ packet := zap.NewPacket(TypeAuthentication, 4+len(*T))
packet = packet.AppendUint32(11)
packet = packet.AppendBytes(*T)
return packet
diff --git a/lib/zap/packets/v3.0/authenticationsaslfinal.go b/lib/zap/packets/v3.0/authenticationsaslfinal.go
index 23b7df2e..e6ba0b9a 100644
--- a/lib/zap/packets/v3.0/authenticationsaslfinal.go
+++ b/lib/zap/packets/v3.0/authenticationsaslfinal.go
@@ -22,7 +22,7 @@ func (T *AuthenticationSASLFinal) ReadFromPacket(packet zap.Packet) bool {
}
func (T *AuthenticationSASLFinal) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthentication)
+ packet := zap.NewPacket(TypeAuthentication, 4+len(*T))
packet = packet.AppendUint32(12)
packet = packet.AppendBytes(*T)
return packet
diff --git a/lib/zap/packets/v3.0/backendkeydata.go b/lib/zap/packets/v3.0/backendkeydata.go
index a14cea1f..a77b8f2b 100644
--- a/lib/zap/packets/v3.0/backendkeydata.go
+++ b/lib/zap/packets/v3.0/backendkeydata.go
@@ -15,7 +15,7 @@ func (T *BackendKeyData) ReadFromPacket(packet zap.Packet) bool {
}
func (T *BackendKeyData) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeBackendKeyData)
+ packet := zap.NewPacket(TypeBackendKeyData, 8)
packet = packet.AppendBytes(T.CancellationKey[:])
return packet
}
diff --git a/lib/zap/packets/v3.0/bind.go b/lib/zap/packets/v3.0/bind.go
index 7eff7ff6..c3eec8aa 100644
--- a/lib/zap/packets/v3.0/bind.go
+++ b/lib/zap/packets/v3.0/bind.go
@@ -52,7 +52,19 @@ func (T *Bind) ReadFromPacket(packet zap.Packet) bool {
}
func (T *Bind) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeBind)
+ size := 0
+ size += len(T.Destination) + 1
+ size += len(T.Source) + 1
+ size += 2
+ size += len(T.ParameterFormatCodes) * 2
+ size += 2
+ for _, v := range T.ParameterValues {
+ size += 4 + len(v)
+ }
+ size += 2
+ size += len(T.ResultFormatCodes) * 2
+
+ packet := zap.NewPacket(TypeBind, size)
packet = packet.AppendString(T.Destination)
packet = packet.AppendString(T.Source)
packet = packet.AppendUint16(uint16(len(T.ParameterFormatCodes)))
diff --git a/lib/zap/packets/v3.0/close.go b/lib/zap/packets/v3.0/close.go
index b8e25a00..66df9786 100644
--- a/lib/zap/packets/v3.0/close.go
+++ b/lib/zap/packets/v3.0/close.go
@@ -17,7 +17,7 @@ func (T *Close) ReadFromPacket(packet zap.Packet) bool {
}
func (T *Close) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeClose)
+ packet := zap.NewPacket(TypeClose, 2+len(T.Target))
packet = packet.AppendUint8(T.Which)
packet = packet.AppendString(T.Target)
return packet
diff --git a/lib/zap/packets/v3.0/copyfail.go b/lib/zap/packets/v3.0/copyfail.go
index 74eab124..7ed87ee5 100644
--- a/lib/zap/packets/v3.0/copyfail.go
+++ b/lib/zap/packets/v3.0/copyfail.go
@@ -15,7 +15,7 @@ func (T *CopyFail) ReadFromPacket(packet zap.Packet) bool {
}
func (T *CopyFail) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeCopyFail)
+ packet := zap.NewPacket(TypeCopyFail, len(T.Reason)+1)
packet = packet.AppendString(T.Reason)
return packet
}
diff --git a/lib/zap/packets/v3.0/describe.go b/lib/zap/packets/v3.0/describe.go
index 4035d8c0..d9276626 100644
--- a/lib/zap/packets/v3.0/describe.go
+++ b/lib/zap/packets/v3.0/describe.go
@@ -17,7 +17,7 @@ func (T *Describe) ReadFromPacket(packet zap.Packet) bool {
}
func (T *Describe) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeDescribe)
+ packet := zap.NewPacket(TypeDescribe, len(T.Target)+2)
packet = packet.AppendUint8(T.Which)
packet = packet.AppendString(T.Target)
return packet
diff --git a/lib/zap/packets/v3.0/errorresponse.go b/lib/zap/packets/v3.0/errorresponse.go
index 6c6a97a6..b707b783 100644
--- a/lib/zap/packets/v3.0/errorresponse.go
+++ b/lib/zap/packets/v3.0/errorresponse.go
@@ -57,7 +57,15 @@ func (T *ErrorResponse) ReadFromPacket(packet zap.Packet) bool {
}
func (T *ErrorResponse) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeErrorResponse)
+ size := 1
+ size += len(T.Error.Severity()) + 2
+ size += len(T.Error.Code()) + 2
+ size += len(T.Error.Message()) + 2
+ for _, field := range T.Error.Extra() {
+ size += len(field.Value) + 2
+ }
+
+ packet := zap.NewPacket(TypeErrorResponse, size)
packet = packet.AppendUint8('S')
packet = packet.AppendString(string(T.Error.Severity()))
diff --git a/lib/zap/packets/v3.0/execute.go b/lib/zap/packets/v3.0/execute.go
index db0525e6..8150bcfe 100644
--- a/lib/zap/packets/v3.0/execute.go
+++ b/lib/zap/packets/v3.0/execute.go
@@ -17,7 +17,7 @@ func (T *Execute) ReadFromPacket(packet zap.Packet) bool {
}
func (T *Execute) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeExecute)
+ packet := zap.NewPacket(TypeExecute, len(T.Target)+5)
packet = packet.AppendString(T.Target)
packet = packet.AppendInt32(T.MaxRows)
return packet
diff --git a/lib/zap/packets/v3.0/negotiateprotocolversion.go b/lib/zap/packets/v3.0/negotiateprotocolversion.go
index 120c23a7..88388b49 100644
--- a/lib/zap/packets/v3.0/negotiateprotocolversion.go
+++ b/lib/zap/packets/v3.0/negotiateprotocolversion.go
@@ -28,7 +28,12 @@ func (T *NegotiateProtocolVersion) ReadFromPacket(packet zap.Packet) bool {
}
func (T *NegotiateProtocolVersion) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeNegotiateProtocolVersion)
+ size := 8
+ for _, v := range T.UnrecognizedOptions {
+ size += len(v) + 1
+ }
+
+ packet := zap.NewPacket(TypeNegotiateProtocolVersion, size)
packet = packet.AppendInt32(T.MinorProtocolVersion)
packet = packet.AppendInt32(int32(len(T.UnrecognizedOptions)))
for _, v := range T.UnrecognizedOptions {
diff --git a/lib/zap/packets/v3.0/parameterstatus.go b/lib/zap/packets/v3.0/parameterstatus.go
index 800790fd..f6dff350 100644
--- a/lib/zap/packets/v3.0/parameterstatus.go
+++ b/lib/zap/packets/v3.0/parameterstatus.go
@@ -17,7 +17,7 @@ func (T *ParameterStatus) ReadFromPacket(packet zap.Packet) bool {
}
func (T *ParameterStatus) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeParameterStatus)
+ packet := zap.NewPacket(TypeParameterStatus, len(T.Key)+len(T.Value)+2)
packet = packet.AppendString(T.Key)
packet = packet.AppendString(T.Value)
return packet
diff --git a/lib/zap/packets/v3.0/parse.go b/lib/zap/packets/v3.0/parse.go
index 4461d6b5..d58c4ea7 100644
--- a/lib/zap/packets/v3.0/parse.go
+++ b/lib/zap/packets/v3.0/parse.go
@@ -27,7 +27,7 @@ func (T *Parse) ReadFromPacket(packet zap.Packet) bool {
}
func (T *Parse) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeParse)
+ packet := zap.NewPacket(TypeParse, len(T.Destination)+len(T.Query)+4+len(T.ParameterDataTypes)*4)
packet = packet.AppendString(T.Destination)
packet = packet.AppendString(T.Query)
packet = packet.AppendInt16(int16(len(T.ParameterDataTypes)))
diff --git a/lib/zap/packets/v3.0/passwordmessage.go b/lib/zap/packets/v3.0/passwordmessage.go
index e6af2d58..087f1959 100644
--- a/lib/zap/packets/v3.0/passwordmessage.go
+++ b/lib/zap/packets/v3.0/passwordmessage.go
@@ -17,7 +17,7 @@ func (T *PasswordMessage) ReadFromPacket(packet zap.Packet) bool {
}
func (T *PasswordMessage) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthenticationResponse)
+ packet := zap.NewPacket(TypeAuthenticationResponse, len(T.Password)+1)
packet = packet.AppendString(T.Password)
return packet
}
diff --git a/lib/zap/packets/v3.0/query.go b/lib/zap/packets/v3.0/query.go
index 8e2ff67d..6cad7ca5 100644
--- a/lib/zap/packets/v3.0/query.go
+++ b/lib/zap/packets/v3.0/query.go
@@ -13,7 +13,7 @@ func (T *Query) ReadFromPacket(packet zap.Packet) bool {
}
func (T *Query) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeQuery)
+ packet := zap.NewPacket(TypeQuery, len(*T)+1)
packet = packet.AppendString(string(*T))
return packet
}
diff --git a/lib/zap/packets/v3.0/readyforquery.go b/lib/zap/packets/v3.0/readyforquery.go
index a40055cf..102e675f 100644
--- a/lib/zap/packets/v3.0/readyforquery.go
+++ b/lib/zap/packets/v3.0/readyforquery.go
@@ -15,7 +15,7 @@ func (T *ReadyForQuery) ReadFromPacket(packet zap.Packet) bool {
}
func (T *ReadyForQuery) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeReadyForQuery)
+ packet := zap.NewPacket(TypeReadyForQuery, 1)
packet = packet.AppendUint8(byte(*T))
return packet
}
diff --git a/lib/zap/packets/v3.0/saslinitialresponse.go b/lib/zap/packets/v3.0/saslinitialresponse.go
index f73eab53..c3f816f3 100644
--- a/lib/zap/packets/v3.0/saslinitialresponse.go
+++ b/lib/zap/packets/v3.0/saslinitialresponse.go
@@ -27,7 +27,7 @@ func (T *SASLInitialResponse) ReadFromPacket(packet zap.Packet) bool {
}
func (T *SASLInitialResponse) IntoPacket() zap.Packet {
- packet := zap.NewPacket(TypeAuthenticationResponse)
+ packet := zap.NewPacket(TypeAuthenticationResponse, len(T.Mechanism)+5+len(T.InitialResponse))
packet = packet.AppendString(T.Mechanism)
packet = packet.AppendInt32(int32(len(T.InitialResponse)))
packet = packet.AppendBytes(T.InitialResponse)
--
GitLab