github.com/glycerine/xcryptossh@v7.0.4+incompatible/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  	"context"
     9  	"io"
    10  	"sync"
    11  	"testing"
    12  )
    13  
    14  // An in-memory packetConn. It is safe to call Close and writePacket
    15  // from different goroutines.
    16  type memTransport struct {
    17  	eof     bool
    18  	pending [][]byte
    19  	write   *memTransport
    20  	sync.Mutex
    21  	*sync.Cond
    22  	idle *IdleTimer
    23  }
    24  
    25  func (t *memTransport) timeout() {
    26  	t.Signal()
    27  }
    28  
    29  func (t *memTransport) readPacket(ctx context.Context) ([]byte, error) {
    30  	t.Lock()
    31  	defer t.Unlock()
    32  	for {
    33  		if len(t.pending) > 0 {
    34  			r := t.pending[0]
    35  			t.pending = t.pending[1:]
    36  			return r, nil
    37  		}
    38  		if t.eof {
    39  			return nil, io.EOF
    40  		}
    41  
    42  		if t.idle != nil {
    43  			select {
    44  			case timedOut := <-t.idle.TimedOut:
    45  				if timedOut != "" {
    46  					return nil, newErrTimeout(timedOut, t.idle)
    47  				}
    48  			case <-t.idle.Halt.ReqStopChan():
    49  				return nil, io.EOF
    50  			}
    51  		}
    52  		t.Cond.Wait()
    53  	}
    54  }
    55  
    56  func (t *memTransport) closeSelf() error {
    57  	t.Lock()
    58  	defer t.Unlock()
    59  	if t.eof {
    60  		return io.EOF
    61  	}
    62  	t.eof = true
    63  	t.Cond.Broadcast()
    64  	return nil
    65  }
    66  
    67  func (t *memTransport) Close() error {
    68  	err := t.write.closeSelf()
    69  	t.closeSelf()
    70  	return err
    71  }
    72  
    73  func (t *memTransport) writePacket(p []byte) error {
    74  	t.write.Lock()
    75  	defer t.write.Unlock()
    76  	if t.write.eof {
    77  		return io.EOF
    78  	}
    79  	c := make([]byte, len(p))
    80  	copy(c, p)
    81  	t.write.pending = append(t.write.pending, c)
    82  	t.write.Cond.Signal()
    83  	return nil
    84  }
    85  
    86  func memPipe() (a, b *memTransport) {
    87  	t1 := memTransport{}
    88  	t2 := memTransport{}
    89  	t1.write = &t2
    90  	t2.write = &t1
    91  	t1.Cond = sync.NewCond(&t1.Mutex)
    92  	t2.Cond = sync.NewCond(&t2.Mutex)
    93  	return &t1, &t2
    94  }
    95  
    96  func TestMemPipe(t *testing.T) {
    97  	defer xtestend(xtestbegin(t))
    98  	a, b := memPipe()
    99  	if err := a.writePacket([]byte{42}); err != nil {
   100  		t.Fatalf("writePacket: %v", err)
   101  	}
   102  	if err := a.Close(); err != nil {
   103  		t.Fatal("Close: ", err)
   104  	}
   105  	ctx := context.Background()
   106  
   107  	p, err := b.readPacket(ctx)
   108  	if err != nil {
   109  		t.Fatal("readPacket: ", err)
   110  	}
   111  	if len(p) != 1 || p[0] != 42 {
   112  		t.Fatalf("got %v, want {42}", p)
   113  	}
   114  	p, err = b.readPacket(ctx)
   115  	if err != io.EOF {
   116  		t.Fatalf("got %v, %v, want EOF", p, err)
   117  	}
   118  }
   119  
   120  func TestDoubleClose(t *testing.T) {
   121  	defer xtestend(xtestbegin(t))
   122  	a, _ := memPipe()
   123  	err := a.Close()
   124  	if err != nil {
   125  		t.Errorf("Close: %v", err)
   126  	}
   127  	err = a.Close()
   128  	if err != io.EOF {
   129  		t.Errorf("expect EOF on double close.")
   130  	}
   131  }