BTP Library

.gitignore

@@ -12,3 +12,6 @@
 
 # Project-local glide cache, RE: https://github.com/Masterminds/glide/issues/736
 .glide/
+
+/apps/newuser/newuser
+/apps/helloworld/helloworld

agent/handlers/executable.go

@@ -9,7 +9,8 @@ import (
 	"net"
 	"os"
 	"os/exec"
-	"path/filepath"
+
+	"github.com/BKellogg/DistributedLoadTester/shared/dir"
 )
 
 // ExecutableHandler handles executables
@@ -29,7 +30,7 @@ func ExecutableHandler(conn net.Conn) error {
 	log.Printf("received connection from %s\n", conn.RemoteAddr().String())
 	write(fmt.Sprintf("Connection recieved. Processing request...\n"), conn)
 
-	fp := currentDir() + "/command"
+	fp := dir.CurrentDir() + "/command"
 
 	// Open or create the file that the bytes will be written to.
 	// Assign the executable permission to the file to the current user.
@@ -115,10 +116,3 @@ func int64FromConn(conn net.Conn) (int64, error) {
 	err := binary.Read(conn, binary.LittleEndian, &size)
 	return size, err
 }
-
-// currentDir returns the directory that the application
-// is being exucuted in
-func currentDir() string {
-	dir, _ := filepath.Abs(filepath.Dir(os.Args[0]))
-	return dir
-}

btp/client.go

@@ -0,0 +1,96 @@
+package btp
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"net"
+	"os"
+)
+
+// RequestBuilder represents a builder
+// for BTP requests
+type RequestBuilder struct {
+	hasPayload  bool
+	beenSent    bool
+	address     string
+	file        string
+	payloadSize int64
+	response    io.ReadCloser
+}
+
+// NewRequestBuilder returns an empty RequestBuilder
+// pointed at the given address.
+func NewRequestBuilder(address string) *RequestBuilder {
+	return &RequestBuilder{hasPayload: false, address: address}
+}
+
+// SetFile sets the given file path as the file
+// to send when the request is sent.
+func (rb *RequestBuilder) SetFile(filepath string) {
+	rb.hasPayload = true
+	rb.file = filepath
+}
+
+// Send sends the RequestBuilder's request.
+// Returns an error if one occurred.
+func (rb *RequestBuilder) Send() error {
+	if !rb.hasPayload {
+		return fmt.Errorf("btp: cannot send a request with no body")
+	}
+	if rb.beenSent {
+		return fmt.Errorf("btp: cannot sent a request that has already been sent")
+	}
+
+	// open the file of this request builder
+	// report any errors that occur.
+	f, err := os.Open(rb.file)
+	if err != nil {
+		return fmt.Errorf("btp: error opening file: %v", err)
+	}
+
+	// obtain the file statistic so we can get the size of the file
+	// report any errors that occur.
+	fstat, err := f.Stat()
+	if err != nil {
+		return fmt.Errorf("btp: error obtaining file statistics: %v", err)
+	}
+
+	// open the connection to this request builder's
+	// address. Report any errors that occur.
+	conn, err := net.Dial("tcp", rb.address)
+	if err != nil {
+		return fmt.Errorf("btp: error dialing address: %v", rb.address)
+	}
+
+	// write the size of the file to the connection
+	// report any errors that occur.
+	if err = binary.Write(conn, binary.LittleEndian, fstat.Size()); err != nil {
+		conn.Close()
+		return fmt.Errorf("btp: error writing file size: %v", err)
+	}
+
+	// copy the contents of the file into the connection
+	// report any errors that occur.
+	if _, err = io.Copy(conn, f); err != nil {
+		conn.Close()
+		return fmt.Errorf("btp: error copying file into connection: %v", err)
+	}
+
+	rb.response = conn
+	rb.beenSent = true
+
+	return nil
+}
+
+// Response gets the response readCloser from the connection.
+// Returns an error if one occurred.
+func (rb *RequestBuilder) Response() (io.ReadCloser, error) {
+	if !rb.beenSent {
+		return nil, fmt.Errorf("btp: cannot get the response of a request that has not been sent")
+	}
+	if rb.response == nil {
+		return nil, fmt.Errorf("btp: response is nil")
+	}
+	return rb.response, nil
+}

btp/handler.go

@@ -0,0 +1,5 @@
+package btp
+
+// Handler defines the type of function that can be
+// used as a BTP Handler
+type Handler func(ResponseWriter, *Request)

btp/reader.go

@@ -0,0 +1,50 @@
+package btp
+
+import (
+	"io"
+)
+
+// Reader returns an io.Reader that
+// returns io.EOF when it reads a
+// specific number of bytes, regardless
+// of if there are more bytes that could be
+// read. The client of this reader will be
+// unable to change the size of the reader
+// so setting it to the size of the body is
+// fine.
+type Reader struct {
+	size   int64
+	reader io.Reader
+}
+
+// Read reads from the current reader into the destination
+// byte slice until either the destination is full, or
+// the number of bytes specified when the reader was
+// created has been hit.
+func (r *Reader) Read(dst []byte) (int, error) {
+	if r.size == 0 {
+		return 0, io.EOF
+	}
+	numBytes, err := r.reader.Read(dst)
+	r.size -= int64(numBytes)
+	return numBytes, err
+}
+
+// newReader creates a new Reader that wraps the given
+// reader. The returned reader will return io.EOF when
+// the source reader is exhausted or the number of bytes
+// read is equal to the given size.
+func newReader(size int64, r io.Reader) *Reader {
+	return &Reader{size: size, reader: r}
+}
+
+// min returns the min value of the given
+// integers. Needed since the math package
+// doesn't implement this for ints, only
+// floats.
+func min(x, y int) int {
+	if x < y {
+		return x
+	}
+	return y
+}

btp/reader_test.go

@@ -0,0 +1,83 @@
+package btp
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"testing"
+)
+
+func TestNewReader(t *testing.T) {
+	reader1 := bytes.NewReader([]byte("some string"))
+	reader2 := ioutil.NopCloser(reader1)
+
+	cases := []struct {
+		size   int64
+		reader io.Reader
+	}{
+		{
+			1,
+			reader1,
+		},
+		{
+			18329320,
+			reader2,
+		},
+	}
+	for _, c := range cases {
+		r := newReader(c.size, c.reader)
+		if r.size != c.size {
+			t.Fatalf("error in NewReader: sizes %d and %d did not match", r.size, c.size)
+		}
+	}
+}
+
+func TestRead(t *testing.T) {
+	cases := []struct {
+		reader        *Reader
+		expectedBytes []byte
+	}{
+		{
+			newReader(10, bytes.NewReader([]byte("hi there!!"))),
+			[]byte("hi there!!"),
+		},
+		{
+			newReader(100, bytes.NewReader([]byte("hi there!!"))),
+			[]byte("hi there!!"),
+		},
+		{
+			newReader(10, bytes.NewReader([]byte("hi there!! how are you doing today?"))),
+			[]byte("hi there!!"),
+		},
+		{
+			newReader(1000, bytes.NewReader([]byte("hi there!! how are you doing today?"))),
+			[]byte("hi there!! how are you doing today?"),
+		},
+		{
+			newReader(0, bytes.NewReader([]byte("hi there!! how are you doing today?"))),
+			[]byte(""),
+		},
+		{
+			newReader(4, bytes.NewReader([]byte("pink fluffy ponies make me happy"))),
+			[]byte("pink"),
+		},
+	}
+	for _, c := range cases {
+		dst := make([]byte, len(c.expectedBytes))
+		numRead, err := c.reader.Read(dst)
+		if err != nil {
+			if err == io.EOF {
+				continue
+			}
+			t.Fatalf("error reading from Reader: %v", err)
+		}
+		if numRead != len(c.expectedBytes) {
+			t.Fatalf("%d bytes read did not match expected number %d", numRead, len(c.expectedBytes))
+		}
+		stringIn, stringOut := string(c.expectedBytes), string(dst)
+		if stringIn != stringOut {
+			t.Fatalf("input %s did not match %s output", stringIn, stringOut)
+		}
+
+	}
+}

btp/server.go

@@ -0,0 +1,45 @@
+package btp
+
+import (
+	"fmt"
+	"log"
+	"net"
+)
+
+// Listen starts a btp server listening on the given
+// address and handles all requests with the given
+// handler. Returns an error if one occurred.
+//
+// Only errors encountered during startup will be returned.
+// Errors encountered during while processing a specific
+// connection during any point in it's lifecycle will not
+// be returned here.
+//
+// This function is a blocking function and will never exit
+// once properly started.
+func Listen(addr string, handler Handler) error {
+	l, err := net.Listen("tcp", addr)
+	if err != nil {
+		return fmt.Errorf("error listening in %s: %v", addr, err)
+	}
+	for {
+		conn, err := l.Accept()
+		if err != nil {
+			log.Printf("error accepting connection: %v\n", err)
+		}
+		go serveRequest(conn, handler)
+	}
+}
+
+// serveRequest serves the given connection with the
+// given handler. Closes the connection when the serving
+// is complete.
+func serveRequest(conn net.Conn, handler Handler) {
+	w, r, err := fullCycleFromConn(conn)
+	if err != nil {
+		log.Printf("error getting conn lifecycle: %v", err)
+		conn.Close()
+	}
+	handler(w, r)
+	conn.Close()
+}