vendor/golang.org/x/crypto/ssh/mempipe_test.go

   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 }