11package prefix
22
33import (
4+ "bufio"
45 "crypto/rand"
56 "fmt"
67 "io"
@@ -11,6 +12,7 @@ import (
1112 "github.com/refraction-networking/conjure/pkg/core"
1213 pb "github.com/refraction-networking/gotapdance/protobuf"
1314 "google.golang.org/protobuf/proto"
15+ "google.golang.org/protobuf/types/known/anypb"
1416)
1517
1618// ClientTransport implements the client side transport interface for the Min transport. The
@@ -42,6 +44,7 @@ type ClientParams struct {
4244// behavior like a rand prefix for example.
4345type Prefix interface {
4446 Bytes () []byte
47+ FlushAfterPrefix () bool
4548 ID () PrefixID
4649 DstPort ([]byte ) uint16
4750}
@@ -91,16 +94,42 @@ func (t *ClientTransport) GetParams() (proto.Message, error) {
9194 return t .parameters , nil
9295}
9396
97+ // ParseParams gives the specific transport an option to parse a generic object into parameters
98+ // provided by the station in the registration response during registration.
99+ func (t ClientTransport ) ParseParams (data * anypb.Any ) (any , error ) {
100+ if data == nil {
101+ return nil , nil
102+ }
103+
104+ var m = & pb.PrefixTransportParams {}
105+ err := transports .UnmarshalAnypbTo (data , m )
106+ return m , err
107+ }
108+
94109// SetParams allows the caller to set parameters associated with the transport, returning an
95110// error if the provided generic message is not compatible or the parameters are otherwise invalid
96- func (t * ClientTransport ) SetParams (p any ) error {
111+ func (t * ClientTransport ) SetParams (p any , unchecked ... bool ) error {
97112 prefixParams , ok := p .(* pb.PrefixTransportParams )
98113 if ! ok {
99- return ErrBadParams
114+ return fmt . Errorf ( "%w, incorrect param type" , ErrBadParams )
100115 }
101116
102117 if prefixParams == nil {
103- return ErrBadParams
118+ return fmt .Errorf ("%w, nil params" , ErrBadParams )
119+ }
120+
121+ if len (unchecked ) != 0 && unchecked [0 ] {
122+ // Overwrite the prefix bytes and type without checking the default set. This is used for
123+ // RegResponse where the registrar may override the chosen prefix with a prefix outside of
124+ // the prefixes that the client known about.
125+ t .parameters = prefixParams
126+ t .Prefix = & clientPrefix {
127+ bytes : prefixParams .GetPrefix (),
128+ id : PrefixID (prefixParams .GetPrefixId ()),
129+ flushAfterPrefix : prefixParams .GetFlushAfterPrefix (),
130+ }
131+
132+ return nil
104133 }
105134
106135 if prefix , ok := DefaultPrefixes [PrefixID (prefixParams .GetPrefixId ())]; ok {
@@ -163,26 +192,6 @@ func (t *ClientTransport) GetDstPort(seed []byte) (uint16, error) {
163192 return t .Prefix .DstPort (seed ), nil
164193}
165194
166- // Build is specific to the Prefix transport, providing a utility function for building the
167- // prefix that the client should write to the wire before sending any client bytes.
168- func (t * ClientTransport ) Build () ([]byte , error ) {
169- if t .Prefix == nil {
170- return nil , ErrBadParams
171- }
172-
173- // Send hmac(seed, str) bytes to indicate to station (min transport)
174- prefix := t .Prefix .Bytes ()
175-
176- if t .TagObfuscator == nil {
177- t .TagObfuscator = transports.CTRObfuscator {}
178- }
179- obfuscatedID , err := t .TagObfuscator .Obfuscate (t .connectTag , t .stationPublicKey [:])
180- if err != nil {
181- return nil , err
182- }
183- return append (prefix , obfuscatedID ... ), nil
184- }
185-
186195// PrepareKeys provides an opportunity for the transport to integrate the station public key
187196// as well as bytes from the deterministic random generator associated with the registration
188197// that this ClientTransport is attached to.
@@ -195,29 +204,52 @@ func (t *ClientTransport) PrepareKeys(pubkey [32]byte, sharedSecret []byte, hkdf
195204// WrapConn gives the transport the opportunity to perform a handshake and wrap / transform the
196205// incoming and outgoing bytes send by the implementing client.
197206func (t * ClientTransport ) WrapConn (conn net.Conn ) (net.Conn , error ) {
198- // Send hmac(seed, str) bytes to indicate to station (min transport) generated during Prepare(...)
199-
200- // // Send hmac(seed, str) bytes to indicate to station (min transport)
201- // connectTag := core.ConjureHMAC(reg.keys.SharedSecret, "PrefixTransportHMACString")
207+ if t .Prefix == nil {
208+ return nil , ErrBadParams
209+ }
202210
203- prefix , err := t .Build ()
204- if err != nil {
205- return nil , fmt .Errorf ("failed to build prefix: %w" , err )
211+ if t .TagObfuscator == nil {
212+ t .TagObfuscator = transports.CTRObfuscator {}
206213 }
207214
208- _ , err = conn . Write ( prefix )
215+ obfuscatedID , err := t . TagObfuscator . Obfuscate ( t . connectTag , t . stationPublicKey [:] )
209216 if err != nil {
210217 return nil , err
211218 }
219+
220+ w := bufio .NewWriter (conn )
221+
222+ var msg []byte = t .Prefix .Bytes ()
223+ if t .Prefix .FlushAfterPrefix () {
224+ if _ , err := w .Write (msg ); err != nil {
225+ return nil , err
226+ }
227+
228+ w .Flush ()
229+ if _ , err := w .Write (obfuscatedID ); err != nil {
230+ return nil , err
231+ }
232+
233+ w .Flush ()
234+ } else {
235+ msg = append (msg , obfuscatedID ... )
236+ if _ , err := w .Write (msg ); err != nil {
237+ return nil , err
238+ }
239+
240+ w .Flush ()
241+ }
242+
212243 return conn , nil
213244}
214245
215246// ---
216247
217248type clientPrefix struct {
218- bytes []byte
219- id PrefixID
220- port uint16
249+ bytes []byte
250+ id PrefixID
251+ port uint16
252+ flushAfterPrefix bool
221253
222254 // // Function allowing encoding / transformation of obfuscated ID bytes after they have been
223255 // // obfuscated. Examples - base64 encode, padding
@@ -242,6 +274,10 @@ func (c *clientPrefix) DstPort([]byte) uint16 {
242274 return c .port
243275}
244276
277+ func (c * clientPrefix ) FlushAfterPrefix () bool {
278+ return c .flushAfterPrefix
279+ }
280+
245281// ---
246282
247283// TryFromID returns a Prefix based on the Prefix ID. This is useful for non-static prefixes like the
0 commit comments