diff --git a/wrappers/packetack/client_abridged.go b/wrappers/packetack/client_abridged.go index c4e68d2..c9b26e1 100644 --- a/wrappers/packetack/client_abridged.go +++ b/wrappers/packetack/client_abridged.go @@ -25,6 +25,7 @@ type wrapperClientAbridged struct { func (w *wrapperClientAbridged) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, error) { buf := bytes.Buffer{} + buf.Grow(1) if _, err := io.CopyN(&buf, w.parent, 1); err != nil { return nil, fmt.Errorf("cannot read message length: %w", err) } @@ -37,6 +38,7 @@ func (w *wrapperClientAbridged) Read(acks *conntypes.ConnectionAcks) (conntypes. } if msgLength == clientAbridgedSmallPacketLength { + buf.Grow(3) if _, err := io.CopyN(&buf, w.parent, 3); err != nil { return nil, fmt.Errorf("cannot read correct message length: %w", err) } @@ -47,6 +49,7 @@ func (w *wrapperClientAbridged) Read(acks *conntypes.ConnectionAcks) (conntypes. msgLength *= 4 buf.Reset() + buf.Grow(int(msgLength)) if _, err := io.CopyN(&buf, w.parent, int64(msgLength)); err != nil { return nil, fmt.Errorf("cannot read message: %w", err) } diff --git a/wrappers/packetack/client_intermediate.go b/wrappers/packetack/client_intermediate.go new file mode 100644 index 0000000..0415c8c --- /dev/null +++ b/wrappers/packetack/client_intermediate.go @@ -0,0 +1,85 @@ +package packetack + +import ( + "bytes" + "encoding/binary" + "fmt" + "io" + "net" + + "go.uber.org/zap" + + "github.com/9seconds/mtg/conntypes" +) + +const clientIntermediateQuickAckLength = 0x80000000 + +type wrapperClientIntermediate struct { + parent conntypes.StreamReadWriteCloser +} + +func (w *wrapperClientIntermediate) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, error) { + buf := bytes.Buffer{} + + buf.Grow(4) + if _, err := io.CopyN(&buf, w.parent, 4); err != nil { + return nil, fmt.Errorf("cannot read message length: %w", err) + } + length := binary.LittleEndian.Uint32(buf.Bytes()) + + if length > clientIntermediateQuickAckLength { + acks.Quick = true + length -= clientIntermediateQuickAckLength + } + + buf.Reset() + buf.Grow(int(length)) + if _, err := io.CopyN(&buf, w.parent, int64(length)); err != nil { + return nil, fmt.Errorf("cannot read the message: %w", err) + } + + return buf.Bytes(), nil +} + +func (w *wrapperClientIntermediate) Write(packet conntypes.Packet, acks *conntypes.ConnectionAcks) error { + if acks.Simple { + if _, err := w.parent.Write(packet); err != nil { + return fmt.Errorf("cannot send simpleacked packet: %w", err) + } + return nil + } + + length := [4]byte{} + binary.LittleEndian.PutUint32(length[:], uint32(len(packet))) + + if _, err := w.parent.Write(append(length[:], packet...)); err != nil { + return fmt.Errorf("cannot send packet: %w", err) + } + return nil +} + +func (w *wrapperClientIntermediate) Close() error { + return w.parent.Close() +} + +func (w *wrapperClientIntermediate) Conn() net.Conn { + return w.parent.Conn() +} + +func (w *wrapperClientIntermediate) LocalAddr() *net.TCPAddr { + return w.parent.LocalAddr() +} + +func (w *wrapperClientIntermediate) RemoteAddr() *net.TCPAddr { + return w.parent.RemoteAddr() +} + +func (w *wrapperClientIntermediate) Logger() *zap.SugaredLogger { + return w.parent.Logger().Named("client-intermediate") +} + +func NewClientIntermediate(parent conntypes.StreamReadWriteCloser) conntypes.PacketAckFullReadWriteCloser { + return &wrapperClientIntermediate{ + parent: parent, + } +} diff --git a/wrappers/packetack/client_intermediate_secure.go b/wrappers/packetack/client_intermediate_secure.go new file mode 100644 index 0000000..f4993e4 --- /dev/null +++ b/wrappers/packetack/client_intermediate_secure.go @@ -0,0 +1,54 @@ +package packetack + +import ( + "bytes" + "encoding/binary" + "fmt" + "math/rand" + + "github.com/9seconds/mtg/conntypes" +) + +type wrapperClientIntermediateSecure struct { + wrapperClientIntermediate +} + +func (w *wrapperClientIntermediateSecure) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, error) { + data, err := w.wrapperClientIntermediate.Read(acks) + if err != nil { + return nil, err + } + length := len(data) - (len(data) % 4) + + return data[:length], nil +} + +func (w *wrapperClientIntermediateSecure) Write(packet conntypes.Packet, acks *conntypes.ConnectionAcks) error { + if acks.Simple { + if _, err := w.parent.Write(packet); err != nil { + return fmt.Errorf("cannot send simpleacked packet: %w", err) + } + return nil + } + + buf := bytes.Buffer{} + paddingLength := rand.Intn(4) + buf.Grow(4 + len(packet) + paddingLength) + + binary.Write(&buf, binary.LittleEndian, uint32(len(packet)+paddingLength)) + buf.Write(packet) + buf.Write(make([]byte, paddingLength)) + + if _, err := w.parent.Write(buf.Bytes()); err != nil { + return fmt.Errorf("cannot send packet: %w", err) + } + return nil +} + +func NewClientIntermediateSecure(parent conntypes.StreamReadWriteCloser) conntypes.PacketAckFullReadWriteCloser { + return &wrapperClientIntermediateSecure{ + wrapperClientIntermediate: wrapperClientIntermediate{ + parent: parent, + }, + } +}