SSH - Fail code for http://codereview.appspot.com/5908048

server_test.go

package main

import (
	"code.google.com/p/go.crypto/ssh"
	"log"
	"io"
	"bytes"
	"sync"
	"time"
	"math/rand"
)

// The deadlock state with the patch and without the patch
// are different.  Thus I will tentatively say the issues
// are probably different.

const (
	ServerNetwork = "127.0.0.1:23000"
	
	// For this to work without entering dead-lock:
	//   With the Server patch to Adjust window msg: 4 or less.
	//   Without patch: 1 or less.
	DataLenMult = 500



	// Chosen as step size after watching ssh packet size.
	DataLen = 16000 * DataLenMult
	//65536
	// 9 280 000
	// 5 151 232
)

func CopyN(dst io.Writer, src io.Reader, n int64) (written int64, err error) {
	buf := make([]byte,  (rand.Intn(30) + 1) * 1024) //32*1024)
	for written < n {
		l := len(buf)
		if d := n - written; d < int64(l) {
			l = int(d)
		}
		nr, er := src.Read(buf[0:l])
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}
func Copy(dst io.Writer, src io.Reader) (written int64, err error) {
	buf := make([]byte, (rand.Intn(30) + 1) * 1024) //32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er == io.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}


func main() {
	startSshServer()

	for i := 0; i < 1; i++ {
		ok := runSshClient()
		if !ok {
			return
		}
	}
	log.Printf("OK")
}

var inServerCopy = false

// First write date to the server in an easily verifiable data pattern.
// Then read back the data and verify with the same pattern.
// Create a timer to catch a hung status.
func runSshClient() (ok bool) {
	conn, err := ssh.Dial("tcp", ServerNetwork, &ssh.ClientConfig{})
	if err != nil {
		log.Fatal(err)
	}

	session, err := conn.NewSession()
	if err != nil {
		log.Fatal(err)
	}

	// Generate large dump and write a simple pattern to.
	var dumpTo = make([]byte, DataLen)
	for i, _ := range dumpTo {
		dumpTo[i] = byte(i % 255)
	}
	var dumpToBuffer = bytes.NewBuffer(dumpTo)

	// To wait for reader to exit.
	// A channel would work here too.
	var wait = &sync.WaitGroup{}
	wait.Add(1)

	// Create a data buffer, size to known length and reset.
	var dumpFromBuffer = bytes.NewBuffer(make([]byte, len(dumpTo)))
	dumpFromBuffer.Reset()

	// Before any data coping, set a timeout in case we hang.
	var inWriterCopy, inReaderCopy bool
	go func() {
		timeout := time.After(time.Second * 10)
		<- timeout
		log.Fatalf("Client/Server communication hung in:: Server: %t, Write: %t, Read: %t", inServerCopy, inWriterCopy, inReaderCopy)
	}()

	// Read back the data from the server.
	go func() {
		defer session.Close()
		serverStdout, err := session.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}

		// Use prints to determine where it gets "stuck."
		inReaderCopy = true
		n, err := CopyN(dumpFromBuffer, serverStdout, DataLen)
		inReaderCopy = false

		if err != nil {
			if err != io.EOF {
				log.Fatal(err)
			}
		}
		if n != int64(len(dumpTo)) {
			log.Fatalf("Expecting different length return.  Was %d, expected %d", n, len(dumpTo))
		}
		log.Printf("--A")
		wait.Done()
	}()

	// Write it all at once and verify length and pattern
	serverStdin, err := session.StdinPipe()
	if err != nil {
		log.Fatal(err)
	}
	
	// Use prints to determine where it gets "stuck."
	inWriterCopy = true
	written, err := Copy(serverStdin, dumpToBuffer)
	inWriterCopy = false
	if err != nil {
		log.Fatal(err)
	}
	if written != DataLen {
		log.Fatalf("Expected %d, Written %d", DataLen, written)
	}

		log.Printf("--B")
	wait.Wait()
		log.Printf("--C")

	// And verify the data.
	var dumpFrom = dumpFromBuffer.Bytes()

	if len(dumpFrom) != len(dumpTo) {
		log.Fatalf("Expecting different length return.  Was %d, expected %d", len(dumpFrom), len(dumpTo))
	}

	n := 0
	for i, _ := range dumpFrom {
		if dumpFrom[i] != byte(i % 255) {
			log.Fatalf("Expected %d, got %d at %d", i % 255, dumpFrom[i], i)
		}
		n++
	}
	return true
}

func startSshServer() {
	config := new(ssh.ServerConfig)
	//config.PublicKeyCallback = keyAuth
	config.NoClientAuth = true

	err := config.SetRSAPrivateKey([]byte(testServerPrivateKey))
	if err != nil {
		log.Fatalf("Failed to parse private key: %s", err.Error())
	}

	listener, err := ssh.Listen("tcp", ServerNetwork, config)
	if err != nil {
		log.Fatalf("Bind error: %s", err)
	}

	log.Printf("Start: %s", ServerNetwork)
	go func() {
		for {
			sConn, err := listener.Accept()

			err = sConn.Handshake()
			if err != nil {
				if err != io.EOF {
					log.Printf("failed to handshake: %s", err)
				}
				return
			}

			go connRun(sConn)
		}
	}()
}

func connRun(sConn *ssh.ServerConn) {
	for {
		channel, err := sConn.Accept()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Fatalf("Channel Accept: %s", err)
		}

		if channel.ChannelType() != "session" {
			channel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		err = channel.Accept()
		if err != nil {
			log.Fatalf("Channel Accept: %s", err)
		}

		go func() {
			defer channel.Close()

			//var payload []byte
			b := make([]byte, 0)

			//read any headers
			i := 0
			for {
				//prevent many headers being sent and processed
				//shell or exec should be the first or second Request
				if i > 5 {
					return
				}
				i++
				_, err := channel.Read(b)
				if err == nil {
					break
				}
				if err != nil {
					log.Print(err)
				}

			}

			inServerCopy = true
			n, err := CopyN(channel, channel, DataLen)
			inServerCopy = false
			if err != nil {
				if err != io.EOF {
					if(err == io.ErrShortWrite) {
						log.Fatalf("short write, wrote %d, expected %d", n, DataLen)
					}
					log.Fatal(err)
				}
			}
		}()
	}
}


const testServerPrivateKey = `-----BEGIN RSA PRIVATE KEY-----
MIIEpAIBAAKCAQEA19lGVsTqIT5iiNYRgnoY1CwkbETW5cq+Rzk5v/kTlf31XpSU
70HVWkbTERECjaYdXM2gGcbb+sxpq6GtXf1M3kVomycqhxwhPv4Cr6Xp4WT/jkFx
9z+FFzpeodGJWjOH6L2H5uX1Cvr9EDdQp9t9/J32/qBFntY8GwoUI/y/1MSTmMiF
tupdMODN064vd3gyMKTwrlQ8tZM6aYuyOPsutLlUY7M5x5FwMDYvnPDSeyT/Iw0z
s3B+NCyqeeMd2T7YzQFnRATj0M7rM5LoSs7DVqVriOEABssFyLj31PboaoLhOKgc
qoM9khkNzr7FHVvi+DhYM2jD0DwvqZLN6NmnLwIDAQABAoIBAQCGVj+kuSFOV1lT
+IclQYA6bM6uY5mroqcSBNegVxCNhWU03BxlW//BE9tA/+kq53vWylMeN9mpGZea
riEMIh25KFGWXqXlOOioH8bkMsqA8S7sBmc7jljyv+0toQ9vCCtJ+sueNPhxQQxH
D2YvUjfzBQ04I9+wn30BByDJ1QA/FoPsunxIOUCcRBE/7jxuLYcpR+JvEF68yYIh
atXRld4W4in7T65YDR8jK1Uj9XAcNeDYNpT/M6oFLx1aPIlkG86aCWRO19S1jLPT
b1ZAKHHxPMCVkSYW0RqvIgLXQOR62D0Zne6/2wtzJkk5UCjkSQ2z7ZzJpMkWgDgN
ifCULFPBAoGBAPoMZ5q1w+zB+knXUD33n1J+niN6TZHJulpf2w5zsW+m2K6Zn62M
MXndXlVAHtk6p02q9kxHdgov34Uo8VpuNjbS1+abGFTI8NZgFo+bsDxJdItemwC4
KJ7L1iz39hRN/ZylMRLz5uTYRGddCkeIHhiG2h7zohH/MaYzUacXEEy3AoGBANz8
e/msleB+iXC0cXKwds26N4hyMdAFE5qAqJXvV3S2W8JZnmU+sS7vPAWMYPlERPk1
D8Q2eXqdPIkAWBhrx4RxD7rNc5qFNcQWEhCIxC9fccluH1y5g2M+4jpMX2CT8Uv+
3z+NoJ5uDTXZTnLCfoZzgZ4nCZVZ+6iU5U1+YXFJAoGBANLPpIV920n/nJmmquMj
orI1R/QXR9Cy56cMC65agezlGOfTYxk5Cfl5Ve+/2IJCfgzwJyjWUsFx7RviEeGw
64o7JoUom1HX+5xxdHPsyZ96OoTJ5RqtKKoApnhRMamau0fWydH1yeOEJd+TRHhc
XStGfhz8QNa1dVFvENczja1vAoGABGWhsd4VPVpHMc7lUvrf4kgKQtTC2PjA4xoc
QJ96hf/642sVE76jl+N6tkGMzGjnVm4P2j+bOy1VvwQavKGoXqJBRd5Apppv727g
/SM7hBXKFc/zH80xKBBgP/i1DR7kdjakCoeu4ngeGywvu2jTS6mQsqzkK+yWbUxJ
I7mYBsECgYB/KNXlTEpXtz/kwWCHFSYA8U74l7zZbVD8ul0e56JDK+lLcJ0tJffk
gqnBycHj6AhEycjda75cs+0zybZvN4x65KZHOGW/O/7OAWEcZP5TPb3zf9ned3Hl
NsZoFj52ponUM6+99A2CmezFCN16c4mbA//luWF+k3VVqR6BpkrhKw==
-----END RSA PRIVATE KEY-----`

const testClientPrivateKey = `-----BEGIN RSA PRIVATE KEY-----
MIIBOwIBAAJBALdGZxkXDAjsYk10ihwU6Id2KeILz1TAJuoq4tOgDWxEEGeTrcld
r/ZwVaFzjWzxaf6zQIJbfaSEAhqD5yo72+sCAwEAAQJBAK8PEVU23Wj8mV0QjwcJ
tZ4GcTUYQL7cF4+ezTCE9a1NrGnCP2RuQkHEKxuTVrxXt+6OF15/1/fuXnxKjmJC
nxkCIQDaXvPPBi0c7vAxGwNY9726x01/dNbHCE0CBtcotobxpwIhANbbQbh3JHVW
2haQh4fAG5mhesZKAGcxTyv4mQ7uMSQdAiAj+4dzMpJWdSzQ+qGHlHMIBvVHLkqB
y2VdEyF7DPCZewIhAI7GOI/6LDIFOvtPo6Bj2nNmyQ1HU6k/LRtNIXi4c9NJAiAr
rrxx26itVhJmcvoUhOjwuzSlP2bE5VHAvkGB352YBg==
-----END RSA PRIVATE KEY-----`