vendor/golang.org/x/crypto/ssh/test/forward_unix_test.go

   1 // Copyright 2012 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 // +build darwin dragonfly freebsd linux netbsd openbsd
   6
   7 package test
   8
   9 import (
  10         "bytes"
  11         "io"
  12         "io/ioutil"
  13         "math/rand"
  14         "net"
  15         "testing"
  16         "time"
  17 )
  18
  19 type closeWriter interface {
  20         CloseWrite() error
  21 }
  22
  23 func testPortForward(t *testing.T, n, listenAddr string) {
  24         server := newServer(t)
  25         defer server.Shutdown()
  26         conn := server.Dial(clientConfig())
  27         defer conn.Close()
  28
  29         sshListener, err := conn.Listen(n, listenAddr)
  30         if err != nil {
  31                 t.Fatal(err)
  32         }
  33
  34         go func() {
  35                 sshConn, err := sshListener.Accept()
  36                 if err != nil {
  37                         t.Fatalf("listen.Accept failed: %v", err)
  38                 }
  39
  40                 _, err = io.Copy(sshConn, sshConn)
  41                 if err != nil && err != io.EOF {
  42                         t.Fatalf("ssh client copy: %v", err)
  43                 }
  44                 sshConn.Close()
  45         }()
  46
  47         forwardedAddr := sshListener.Addr().String()
  48         netConn, err := net.Dial(n, forwardedAddr)
  49         if err != nil {
  50                 t.Fatalf("net dial failed: %v", err)
  51         }
  52
  53         readChan := make(chan []byte)
  54         go func() {
  55                 data, _ := ioutil.ReadAll(netConn)
  56                 readChan <- data
  57         }()
  58
  59         // Invent some data.
  60         data := make([]byte, 100*1000)
  61         for i := range data {
  62                 data[i] = byte(i % 255)
  63         }
  64
  65         var sent []byte
  66         for len(sent) < 1000*1000 {
  67                 // Send random sized chunks
  68                 m := rand.Intn(len(data))
  69                 n, err := netConn.Write(data[:m])
  70                 if err != nil {
  71                         break
  72                 }
  73                 sent = append(sent, data[:n]...)
  74         }
  75         if err := netConn.(closeWriter).CloseWrite(); err != nil {
  76                 t.Errorf("netConn.CloseWrite: %v", err)
  77         }
  78
  79         read := <-readChan
  80
  81         if len(sent) != len(read) {
  82                 t.Fatalf("got %d bytes, want %d", len(read), len(sent))
  83         }
  84         if bytes.Compare(sent, read) != 0 {
  85                 t.Fatalf("read back data does not match")
  86         }
  87
  88         if err := sshListener.Close(); err != nil {
  89                 t.Fatalf("sshListener.Close: %v", err)
  90         }
  91
  92         // Check that the forward disappeared.
  93         netConn, err = net.Dial(n, forwardedAddr)
  94         if err == nil {
  95                 netConn.Close()
  96                 t.Errorf("still listening to %s after closing", forwardedAddr)
  97         }
  98 }
  99
 100 func TestPortForwardTCP(t *testing.T) {
 101         testPortForward(t, "tcp", "localhost:0")
 102 }
 103
 104 func TestPortForwardUnix(t *testing.T) {
 105         addr, cleanup := newTempSocket(t)
 106         defer cleanup()
 107         testPortForward(t, "unix", addr)
 108 }
 109
 110 func testAcceptClose(t *testing.T, n, listenAddr string) {
 111         server := newServer(t)
 112         defer server.Shutdown()
 113         conn := server.Dial(clientConfig())
 114
 115         sshListener, err := conn.Listen(n, listenAddr)
 116         if err != nil {
 117                 t.Fatal(err)
 118         }
 119
 120         quit := make(chan error, 1)
 121         go func() {
 122                 for {
 123                         c, err := sshListener.Accept()
 124                         if err != nil {
 125                                 quit <- err
 126                                 break
 127                         }
 128                         c.Close()
 129                 }
 130         }()
 131         sshListener.Close()
 132
 133         select {
 134         case <-time.After(1 * time.Second):
 135                 t.Errorf("timeout: listener did not close.")
 136         case err := <-quit:
 137                 t.Logf("quit as expected (error %v)", err)
 138         }
 139 }
 140
 141 func TestAcceptCloseTCP(t *testing.T) {
 142         testAcceptClose(t, "tcp", "localhost:0")
 143 }
 144
 145 func TestAcceptCloseUnix(t *testing.T) {
 146         addr, cleanup := newTempSocket(t)
 147         defer cleanup()
 148         testAcceptClose(t, "unix", addr)
 149 }
 150
 151 // Check that listeners exit if the underlying client transport dies.
 152 func testPortForwardConnectionClose(t *testing.T, n, listenAddr string) {
 153         server := newServer(t)
 154         defer server.Shutdown()
 155         conn := server.Dial(clientConfig())
 156
 157         sshListener, err := conn.Listen(n, listenAddr)
 158         if err != nil {
 159                 t.Fatal(err)
 160         }
 161
 162         quit := make(chan error, 1)
 163         go func() {
 164                 for {
 165                         c, err := sshListener.Accept()
 166                         if err != nil {
 167                                 quit <- err
 168                                 break
 169                         }
 170                         c.Close()
 171                 }
 172         }()
 173
 174         // It would be even nicer if we closed the server side, but it
 175         // is more involved as the fd for that side is dup()ed.
 176         server.clientConn.Close()
 177
 178         select {
 179         case <-time.After(1 * time.Second):
 180                 t.Errorf("timeout: listener did not close.")
 181         case err := <-quit:
 182                 t.Logf("quit as expected (error %v)", err)
 183         }
 184 }
 185
 186 func TestPortForwardConnectionCloseTCP(t *testing.T) {
 187         testPortForwardConnectionClose(t, "tcp", "localhost:0")
 188 }
 189
 190 func TestPortForwardConnectionCloseUnix(t *testing.T) {
 191         addr, cleanup := newTempSocket(t)
 192         defer cleanup()
 193         testPortForwardConnectionClose(t, "unix", addr)
 194 }