pty.go

pty.go

package main

import (
	"os"
	"os/exec"
	"fmt"
	"net"
	"io"
	"time"

	"github.com/coreos/go-etcd/etcd"
)

const (
	connHost = "0.0.0.0"
	connType = "tcp"
	connPort = "3333"
	timeout time.Duration = 10 * time.Second
	etcdTTL time.Duration = timeout * 2
)

// Handles incoming connection
func handleRequest(conn net.Conn, command string) {
	fmt.Println("Executing: ", command)

	cmd := exec.Command("bash", "-c", command)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		panic(err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		panic(err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		panic(err)
	}
	cmd.Start()

	connInCh := make(chan []byte)
	connInErroCh := make(chan error)
	stdoutCh := make(chan []byte)
	stdoutErroCh := make(chan error)
	stderrCh := make(chan []byte)
	stderrErroCh := make(chan error)

	// Setup connection channel
	go readInput(conn, connInCh, connInErroCh)
	// Setup stdout channel
	go readInput(stdout, stdoutCh, stdoutErroCh)
	// Setup stderr channel
	go readInput(stderr, stderrCh, stderrErroCh)

	// Pipe in- and outputs
	loop:
	for {
		select {
			// Received data from connection
			case data := <-connInCh:
				stdin.Write(data)
			// Received an error from connection
			case <-connInErroCh:
				break loop;
			// Received data from stdout
			case data := <-stdoutCh:
				conn.Write(data)
			// Received an error from stdout
			case <-stdoutErroCh:
				break loop;
			// Received data from stdout
			case data := <-stderrCh:
				conn.Write(data)
			// Received an error from stdout
			case <-stderrErroCh:
				break loop;
		}
	}

	// Determine exit code
	err = cmd.Wait();
	statusCode := 0;
	if err != nil {
		statusCode = 1;
	}
	statusResponse := []byte(fmt.Sprintf("%v", statusCode));
	conn.Write(statusResponse)

	conn.Close()
}

func readInput(r io.Reader, ch chan []byte, eCh chan error) {
	for {
		data := make([]byte, 512)
		_, err := r.Read(data)
		if err != nil {
			eCh<- err
			return
		}
		ch<- data
	}
}

func setEtcd(client *etcd.Client, key, value string, ttl uint64) {
	_, err := client.Set(key, value, ttl)
	if err != nil {
		fmt.Println(err)
	}
	fmt.Println("set", key, "->", value)
}

func announcePty(etcdClient *etcd.Client, ttl time.Duration, host, port, runID string) {
	key := "/deis/pty/" + runID
	value := host + ":" + port
	for {
		go setEtcd(etcdClient, key, value, uint64(ttl.Seconds()))
		time.Sleep(timeout)
	}
}

func main() {
	command := os.Getenv("COMMAND")
	host := os.Getenv("HOST")
	externalPort := os.Getenv("EXTERNAL_PORT")
	etcdPort := os.Getenv("ETCD_PORT")
	if etcdPort == "" {
		etcdPort = "4001"
	}
	runID := os.Getenv("RUN_ID")

	// Listen for incoming connections.
	l, err := net.Listen(connType, connHost + ":" + connPort)
	if err != nil {
		fmt.Println("Error listening:", err.Error())
		os.Exit(1)
	}
	// Close the listener when the application closes.
	defer l.Close()

	fmt.Println("Listening on " + connHost + ":" + connPort)

	// Register the pty in etcd
	etcdClient := etcd.NewClient([]string{"http://" + host + ":" + etcdPort})
	go announcePty(etcdClient, etcdTTL, host, externalPort, runID)

	// Listen for an incoming connection.
	conn, err := l.Accept()
	if err != nil {
		fmt.Println("Error accepting: ", err.Error())
		os.Exit(1)
	}

	// Handle connection
	handleRequest(conn, command)
}