Browse Source

lib/tlsutil: Allocate UnionedConnection in one go (#7912)

greatroar 4 years ago
parent
commit
a4489dec30
2 changed files with 19 additions and 18 deletions
  1. 8 6
      lib/tlsutil/tlsutil.go
  2. 11 12
      lib/tlsutil/tlsutil_test.go

+ 8 - 6
lib/tlsutil/tlsutil.go

@@ -185,9 +185,10 @@ func (l *DowngradingListener) AcceptNoWrapTLS() (net.Conn, bool, error) {
 		return nil, false, err
 	}
 
-	var first [1]byte
+	union := &UnionedConnection{Conn: conn}
+
 	conn.SetReadDeadline(time.Now().Add(1 * time.Second))
-	n, err := conn.Read(first[:])
+	n, err := conn.Read(union.first[:])
 	conn.SetReadDeadline(time.Time{})
 	if err != nil || n == 0 {
 		// We hit a read error here, but the Accept() call succeeded so we must not return an error.
@@ -196,22 +197,23 @@ func (l *DowngradingListener) AcceptNoWrapTLS() (net.Conn, bool, error) {
 		return conn, false, ErrIdentificationFailed
 	}
 
-	return &UnionedConnection{&first, conn}, first[0] == 0x16, nil
+	return union, union.first[0] == 0x16, nil
 }
 
 type UnionedConnection struct {
-	first *[1]byte
+	first     [1]byte
+	firstDone bool
 	net.Conn
 }
 
 func (c *UnionedConnection) Read(b []byte) (n int, err error) {
-	if c.first != nil {
+	if !c.firstDone {
 		if len(b) == 0 {
 			// this probably doesn't happen, but handle it anyway
 			return 0, nil
 		}
 		b[0] = c.first[0]
-		c.first = nil
+		c.firstDone = true
 		return 1, nil
 	}
 	return c.Conn.Read(b)

+ 11 - 12
lib/tlsutil/tlsutil_test.go

@@ -60,9 +60,8 @@ func TestUnionedConnection(t *testing.T) {
 				if n != 1 {
 					t.Errorf("%d: first read returned %d bytes, not 1", i, n)
 				}
-				// Check that we've nilled out the "first" thing
-				if conn.(*UnionedConnection).first != nil {
-					t.Errorf("%d: expected first read to clear out the `first` attribute", i)
+				if !conn.(*UnionedConnection).firstDone {
+					t.Errorf("%d: expected first read to set the `firstDone` attribute", i)
 				}
 			}
 			bs = append(bs, buf[:n]...)
@@ -130,8 +129,8 @@ func (f *fakeAccepter) Accept() (net.Conn, error) {
 	return &fakeConn{f.data}, nil
 }
 
-func (f *fakeAccepter) Addr() net.Addr { return nil }
-func (f *fakeAccepter) Close() error   { return nil }
+func (*fakeAccepter) Addr() net.Addr { return nil }
+func (*fakeAccepter) Close() error   { return nil }
 
 type fakeConn struct {
 	data []byte
@@ -146,13 +145,13 @@ func (f *fakeConn) Read(b []byte) (int, error) {
 	return n, nil
 }
 
-func (f *fakeConn) Write(b []byte) (int, error) {
+func (*fakeConn) Write(b []byte) (int, error) {
 	return len(b), nil
 }
 
-func (f *fakeConn) Close() error                     { return nil }
-func (f *fakeConn) LocalAddr() net.Addr              { return nil }
-func (f *fakeConn) RemoteAddr() net.Addr             { return nil }
-func (f *fakeConn) SetDeadline(time.Time) error      { return nil }
-func (f *fakeConn) SetReadDeadline(time.Time) error  { return nil }
-func (f *fakeConn) SetWriteDeadline(time.Time) error { return nil }
+func (*fakeConn) Close() error                     { return nil }
+func (*fakeConn) LocalAddr() net.Addr              { return nil }
+func (*fakeConn) RemoteAddr() net.Addr             { return nil }
+func (*fakeConn) SetDeadline(time.Time) error      { return nil }
+func (*fakeConn) SetReadDeadline(time.Time) error  { return nil }
+func (*fakeConn) SetWriteDeadline(time.Time) error { return nil }