main.go

//chisel end-to-end test
//======================
//
//                    (direct)
//         .--------------->----------------.
//        /    chisel         chisel         \
// request--->client:2001--->server:2002---->fileserver:3000
//        \                                  /
//         '--> crowbar:4001--->crowbar:4002'
//              client           server
//
// crowbar and chisel binaries should be in your PATH

package main

import (
  "flag"
  "fmt"
  "io"
  "io/ioutil"
  "log"
  "net/http"
  "os"
  "os/exec"
  "path"
  "strconv"

  "github.com/jpillora/chisel/share"

  "time"
)

const ENABLE_CROWBAR = false

const (
  B  = 1
  KB = 1000 * B
  MB = 1000 * KB
  GB = 1000 * MB
)

func run() {
  flag.Parse()
  args := flag.Args()
  if len(args) == 0 {
    fatal("go run main.go [test] or [bench]")
  }
  for _, a := range args {
    switch a {
    case "test":
      test()
    case "bench":
      bench()
    }
  }
}

//test
func test() {
  testTunnel("2001", 500)
  testTunnel("2001", 50000)
}

//benchmark
func bench() {
  benchSizes("3000")
  benchSizes("2001")
  if ENABLE_CROWBAR {
    benchSizes("4001")
  }
}

func benchSizes(port string) {
  for size := 1; size <= 100*MB; size *= 10 {
    testTunnel(port, size)
  }
}

func testTunnel(port string, size int) {
  t0 := time.Now()
  resp, err := requestFile(port, size)
  if err != nil {
    fatal(err)
  }
  if resp.StatusCode != 200 {
    fatal(err)
  }

  n, err := io.Copy(ioutil.Discard, resp.Body)
  if err != nil {
    fatal(err)
  }
  t1 := time.Now()
  fmt.Printf(":%s => %d bytes in %s\n", port, size, t1.Sub(t0))
  if int(n) != size {
    fatalf("%d bytes expected, got %d", size, n)
  }
}

//============================

func requestFile(port string, size int) (*http.Response, error) {
  url := "http://127.0.0.1:" + port + "/" + strconv.Itoa(size)
  // fmt.Println(url)
  return http.Get(url)
}

func makeFileServer() *chshare.HTTPServer {
  bsize := 3 * MB
  bytes := make([]byte, bsize)
  //filling huge buffer
  for i := 0; i < len(bytes); i++ {
    bytes[i] = byte(i)
  }

  s := chshare.NewHTTPServer()
  s.SetKeepAlivesEnabled(false)
  handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    rsize, _ := strconv.Atoi(r.URL.Path[1:])
    for rsize >= bsize {
      w.Write(bytes)
      rsize -= bsize
    }
    w.Write(bytes[:rsize])
  })
  s.GoListenAndServe("0.0.0.0:3000", handler)
  return s
}

//============================

func fatal(args ...interface{}) {
  panic(fmt.Sprint(args...))
}
func fatalf(f string, args ...interface{}) {
  panic(fmt.Sprintf(f, args...))
}

//global setup
func main() {

  fs := makeFileServer()
  go func() {
    err := fs.Wait()
    if err != nil {
      fmt.Printf("fs server closed (%s)\n", err)
    }
  }()

  if ENABLE_CROWBAR {
    dir, _ := os.Getwd()
    cd := exec.Command("crowbard",
      `-listen`, "0.0.0.0:4002",
      `-userfile`, path.Join(dir, "userfile"))
    if err := cd.Start(); err != nil {
      fatal(err)
    }
    go func() {
      fatalf("crowbard: %v", cd.Wait())
    }()
    defer cd.Process.Kill()

    time.Sleep(100 * time.Millisecond)

    cf := exec.Command("crowbar-forward",
      "-local=0.0.0.0:4001",
      "-server=http://127.0.0.1:4002",
      "-remote=127.0.0.1:3000",
      "-username", "foo",
      "-password", "bar")
    if err := cf.Start(); err != nil {
      fatal(err)
    }
    defer cf.Process.Kill()
  }

  time.Sleep(100 * time.Millisecond)

  hd := exec.Command("chisel", "server",
    // "-v",
    "--key", "foobar",
    "--port", "2002")
  hd.Stdout = os.Stdout
  if err := hd.Start(); err != nil {
    fatal(err)
  }
  defer hd.Process.Kill()

  time.Sleep(100 * time.Millisecond)

  hf := exec.Command("chisel", "client",
    // "-v",
    "--fingerprint", "ed:f2:cf:3c:56",
    "127.0.0.1:2002",
    "2001:3000")
  hf.Stdout = os.Stdout
  if err := hf.Start(); err != nil {
    fatal(err)
  }
  defer hf.Process.Kill()

  time.Sleep(100 * time.Millisecond)

  defer func() {
    if r := recover(); r != nil {
      log.Print(r)
    }
  }()
  run()

  fs.Close()
}