github.com/Psiphon-Labs/psiphon-tunnel-core@v2.0.28+incompatible/psiphon/common/crypto/ssh/mempipe_test.go (about)

     1  // Copyright 2013 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package ssh
     6  
     7  import (
     8  	"io"
     9  	"sync"
    10  	"testing"
    11  )
    12  
    13  // An in-memory packetConn. It is safe to call Close and writePacket
    14  // from different goroutines.
    15  type memTransport struct {
    16  	eof     bool
    17  	pending [][]byte
    18  	write   *memTransport
    19  	sync.Mutex
    20  	*sync.Cond
    21  }
    22  
    23  func (t *memTransport) readPacket() ([]byte, error) {
    24  	t.Lock()
    25  	defer t.Unlock()
    26  	for {
    27  		if len(t.pending) > 0 {
    28  			r := t.pending[0]
    29  			t.pending = t.pending[1:]
    30  			return r, nil
    31  		}
    32  		if t.eof {
    33  			return nil, io.EOF
    34  		}
    35  		t.Cond.Wait()
    36  	}
    37  }
    38  
    39  func (t *memTransport) closeSelf() error {
    40  	t.Lock()
    41  	defer t.Unlock()
    42  	if t.eof {
    43  		return io.EOF
    44  	}
    45  	t.eof = true
    46  	t.Cond.Broadcast()
    47  	return nil
    48  }
    49  
    50  func (t *memTransport) Close() error {
    51  	err := t.write.closeSelf()
    52  	t.closeSelf()
    53  	return err
    54  }
    55  
    56  func (t *memTransport) writePacket(p []byte) error {
    57  	t.write.Lock()
    58  	defer t.write.Unlock()
    59  	if t.write.eof {
    60  		return io.EOF
    61  	}
    62  	c := make([]byte, len(p))
    63  	copy(c, p)
    64  	t.write.pending = append(t.write.pending, c)
    65  	t.write.Cond.Signal()
    66  	return nil
    67  }
    68  
    69  func memPipe() (a, b packetConn) {
    70  	t1 := memTransport{}
    71  	t2 := memTransport{}
    72  	t1.write = &t2
    73  	t2.write = &t1
    74  	t1.Cond = sync.NewCond(&t1.Mutex)
    75  	t2.Cond = sync.NewCond(&t2.Mutex)
    76  	return &t1, &t2
    77  }
    78  
    79  func TestMemPipe(t *testing.T) {
    80  	a, b := memPipe()
    81  	if err := a.writePacket([]byte{42}); err != nil {
    82  		t.Fatalf("writePacket: %v", err)
    83  	}
    84  	if err := a.Close(); err != nil {
    85  		t.Fatal("Close: ", err)
    86  	}
    87  	p, err := b.readPacket()
    88  	if err != nil {
    89  		t.Fatal("readPacket: ", err)
    90  	}
    91  	if len(p) != 1 || p[0] != 42 {
    92  		t.Fatalf("got %v, want {42}", p)
    93  	}
    94  	p, err = b.readPacket()
    95  	if err != io.EOF {
    96  		t.Fatalf("got %v, %v, want EOF", p, err)
    97  	}
    98  }
    99  
   100  func TestDoubleClose(t *testing.T) {
   101  	a, _ := memPipe()
   102  	err := a.Close()
   103  	if err != nil {
   104  		t.Errorf("Close: %v", err)
   105  	}
   106  	err = a.Close()
   107  	if err != io.EOF {
   108  		t.Errorf("expect EOF on double close.")
   109  	}
   110  }