Skip to content

Instantly share code, notes, and snippets.

@johanneswuerbach
Created October 14, 2014 05:54
Show Gist options
  • Save johanneswuerbach/6cbc88e3aea3f998a63f to your computer and use it in GitHub Desktop.
Save johanneswuerbach/6cbc88e3aea3f998a63f to your computer and use it in GitHub Desktop.
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)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment