func init() {
conf = config.New()
conf.Username = "******"
conf.KontrolURL = "http://localhost:4099/kite"
conf.KontrolKey = testkeys.Public
conf.KontrolUser = "******"
conf.KiteKey = testutil.NewKiteKey().Raw
conf.ReadEnvironmentVariables()
kontrol.DefaultPort = 4099
kon := kontrol.New(conf.Copy(), "0.1.0")
switch os.Getenv("KONTROL_STORAGE") {
case "etcd":
kon.SetStorage(kontrol.NewEtcd(nil, kon.Kite.Log))
case "postgres":
p := kontrol.NewPostgres(nil, kon.Kite.Log)
kon.SetStorage(p)
kon.SetKeyPairStorage(p)
default:
kon.SetStorage(kontrol.NewEtcd(nil, kon.Kite.Log))
}
kon.AddKeyPair("", testkeys.Public, testkeys.Private)
go kon.Run()
<-kon.Kite.ServerReadyNotify()
}
func TestProxy(t *testing.T) {
conf := config.New()
conf.Username = "******"
conf.KontrolURL = "ws://localhost:6666/kite"
conf.KontrolKey = testkeys.Public
conf.KontrolUser = "******"
conf.KiteKey = testutil.NewKiteKey().Raw
conf.Transport = config.WebSocket // tunnel only works via WebSocket
// start kontrol
color.Green("Starting kontrol")
kontrol.DefaultPort = 6666
kon := kontrol.New(conf.Copy(), "0.1.0")
switch os.Getenv("KONTROL_STORAGE") {
case "etcd":
kon.SetStorage(kontrol.NewEtcd(nil, kon.Kite.Log))
case "postgres":
p := kontrol.NewPostgres(nil, kon.Kite.Log)
kon.SetStorage(p)
kon.SetKeyPairStorage(p)
default:
kon.SetStorage(kontrol.NewEtcd(nil, kon.Kite.Log))
}
kon.AddKeyPair("", testkeys.Public, testkeys.Private)
go kon.Run()
<-kon.Kite.ServerReadyNotify()
DefaultPort = 4999
DefaultPublicHost = "localhost:4999"
prxConf := conf.Copy()
prxConf.DisableAuthentication = true // no kontrol running in test
prx := New(prxConf, "0.1.0", testkeys.Public, testkeys.Private)
prx.Start()
log.Println("Proxy started")
// Proxy kite is ready.
kite1 := kite.New("kite1", "1.0.0")
kite1.Config = conf.Copy()
kite1.HandleFunc("foo", func(r *kite.Request) (interface{}, error) {
return "bar", nil
})
prxClt := kite1.NewClient("http://localhost:4999/kite")
err := prxClt.Dial()
if err != nil {
t.Fatal(err)
}
// Kite1 is connected to proxy.
result, err := prxClt.TellWithTimeout("register", 4*time.Second)
if err != nil {
t.Fatal(err)
}
proxyURL := result.MustString()
log.Printf("Registered to proxy with URL: %s", proxyURL)
if !strings.Contains(proxyURL, "/proxy") {
t.Fatalf("Invalid proxy URL: %s", proxyURL)
}
kite2 := kite.New("kite2", "1.0.0")
kite2.Config = conf.Copy()
kite1remote := kite2.NewClient(proxyURL)
err = kite1remote.Dial()
if err != nil {
t.Fatal(err)
}
// kite2 is connected to kite1 via proxy kite.
result, err = kite1remote.TellWithTimeout("foo", 4*time.Second)
if err != nil {
t.Fatal(err)
}
s := result.MustString()
if s != "bar" {
t.Fatalf("Wrong reply: %s", s)
}
}
func New(c *Config) *kontrol.Kontrol {
modelhelper.Initialize(c.MongoURL)
publicKey, err := ioutil.ReadFile(FindPath(c.PublicKey))
if err != nil {
log.Fatalln(err.Error())
}
privateKey, err := ioutil.ReadFile(FindPath(c.PrivateKey))
if err != nil {
log.Fatalln(err.Error())
}
kiteConf := config.MustGet()
if c.Environment != "" {
kiteConf.Environment = c.Environment
}
if c.Region != "" {
kiteConf.Region = c.Region
}
if c.Port != 0 {
kiteConf.Port = c.Port
}
// TODO: Move the metrics instance somewhere meaningful
met := common.MustInitMetrics(Name)
kon := kontrol.NewWithoutHandlers(kiteConf, Version)
kon.TokenNoNBF = true
kon.Kite.HandleFunc("register",
metricKiteHandler(met, "HandleRegister", kon.HandleRegister),
)
kon.Kite.HandleFunc("registerMachine",
metricKiteHandler(met, "HandleMachine", kon.HandleMachine),
).DisableAuthentication()
kon.Kite.HandleFunc("getKodingKites",
metricKiteHandler(
met, "HandleGetKodingKites", HandleGetKodingKites(kon.HandleGetKites, kiteConf.Environment),
),
)
kon.Kite.HandleFunc("getKites",
metricKiteHandler(met, "HandleGetKites", kon.HandleGetKites),
)
kon.Kite.HandleFunc("getToken",
metricKiteHandler(met, "HandleGetToken", kon.HandleGetToken),
)
kon.Kite.HandleFunc("getKey",
metricKiteHandler(met, "HandleGetKey", kon.HandleGetKey),
)
kon.Kite.HandleHTTPFunc("/heartbeat",
metricHandler(met, "HandleHeartbeat", kon.HandleHeartbeat),
)
kon.Kite.HandleHTTP("/register", throttledHandler(
metricHandler(met, "HandleRegisterHTTP", kon.HandleRegisterHTTP),
))
kon.AddAuthenticator("sessionID", authenticateFromSessionID)
kon.MachineAuthenticate = authenticateMachine
switch c.Storage {
case "etcd":
kon.SetStorage(kontrol.NewEtcd(c.Machines, kon.Kite.Log))
case "postgres":
postgresConf := &kontrol.PostgresConfig{
Host: c.Postgres.Host,
Port: c.Postgres.Port,
Username: c.Postgres.Username,
Password: c.Postgres.Password,
DBName: c.Postgres.DBName,
}
p := kontrol.NewPostgres(postgresConf, kon.Kite.Log)
p.DB.SetMaxOpenConns(20)
kon.SetStorage(p)
s := kontrol.NewCachedStorage(
p,
kontrol.NewMemKeyPairStorageTTL(time.Minute*5),
)
kon.SetKeyPairStorage(s)
// kon.MachineKeyPicker = newMachineKeyPicker(p)
default:
panic(fmt.Sprintf("storage is not found: %q", c.Storage))
}
kon.AddKeyPair("", string(publicKey), string(privateKey))
if c.TLSKeyFile != "" && c.TLSCertFile != "" {
kon.Kite.UseTLSFile(c.TLSCertFile, c.TLSKeyFile)
}
return kon
}
func main() {
conf := new(Kontrol)
multiconfig.New().MustLoad(conf)
publicKey, err := ioutil.ReadFile(conf.PublicKeyFile)
if err != nil {
log.Fatalf("cannot read public key file: %s", err.Error())
}
privateKey, err := ioutil.ReadFile(conf.PrivateKeyFile)
if err != nil {
log.Fatalf("cannot read private key file: %s", err.Error())
}
if conf.Initial {
initialKey(conf, publicKey, privateKey)
os.Exit(0)
}
kiteConf := config.MustGet()
kiteConf.IP = conf.Ip
kiteConf.Port = conf.Port
k := kontrol.New(kiteConf, conf.Version)
if conf.TLSCertFile != "" || conf.TLSKeyFile != "" {
cert, err := tls.LoadX509KeyPair(conf.TLSCertFile, conf.TLSKeyFile)
if err != nil {
log.Fatalf("cannot load TLS certificate: %s", err.Error())
}
k.Kite.TLSConfig = &tls.Config{Certificates: []tls.Certificate{cert}}
}
if conf.RegisterUrl != "" {
k.RegisterURL = conf.RegisterUrl
}
switch os.Getenv("KONTROL_STORAGE") {
case "etcd":
k.SetStorage(kontrol.NewEtcd(conf.Machines, k.Kite.Log))
case "postgres":
postgresConf := &kontrol.PostgresConfig{
Host: conf.Postgres.Host,
Port: conf.Postgres.Port,
Username: conf.Postgres.Username,
Password: conf.Postgres.Password,
DBName: conf.Postgres.DBName,
}
p := kontrol.NewPostgres(postgresConf, k.Kite.Log)
k.SetStorage(p)
k.SetKeyPairStorage(p)
default:
k.SetStorage(kontrol.NewEtcd(conf.Machines, k.Kite.Log))
}
k.AddKeyPair("", string(publicKey), string(privateKey))
k.Kite.SetLogLevel(kite.DEBUG)
k.Run()
}
func TestWebSocketProxy(t *testing.T) {
color.Blue("====> Starting WebSocket test")
conf := config.New()
conf.Username = "******"
conf.KontrolURL = "http://localhost:5555/kite"
conf.KontrolKey = testkeys.Public
conf.KontrolUser = "******"
conf.KiteKey = testutil.NewKiteKey().Raw
conf.ReadEnvironmentVariables()
// start kontrol
color.Green("Starting kontrol")
kontrol.DefaultPort = 5555
kon := kontrol.New(conf.Copy(), "0.1.0")
switch os.Getenv("KONTROL_STORAGE") {
case "etcd":
kon.SetStorage(kontrol.NewEtcd(nil, kon.Kite.Log))
case "postgres":
p := kontrol.NewPostgres(nil, kon.Kite.Log)
kon.SetStorage(p)
kon.SetKeyPairStorage(p)
default:
kon.SetStorage(kontrol.NewEtcd(nil, kon.Kite.Log))
}
kon.AddKeyPair("", testkeys.Public, testkeys.Private)
go kon.Run()
<-kon.Kite.ServerReadyNotify()
// start proxy
color.Green("Starting Proxy and registering to Kontrol")
proxyConf := conf.Copy()
proxyConf.Port = 4999
proxy := New(proxyConf)
proxy.PublicHost = "localhost"
proxy.PublicPort = proxyConf.Port
proxy.Scheme = "http"
go proxy.Run()
<-proxy.ReadyNotify()
proxyRegisterURL := &url.URL{
Scheme: proxy.Scheme,
Host: proxy.PublicHost + ":" + strconv.Itoa(proxy.PublicPort),
Path: "/kite",
}
fmt.Printf("proxyRegisterURL %+v\n", proxyRegisterURL)
_, err := proxy.Kite.Register(proxyRegisterURL)
if err != nil {
t.Error(err)
}
// start now backend kite
color.Green("Starting BackendKite")
backendKite := kite.New("backendKite", "1.0.0")
backendKite.Config = conf.Copy()
backendKite.HandleFunc("foo", func(r *kite.Request) (interface{}, error) {
return "bar", nil
})
backendKite.Config.Port = 7777
kiteUrl := &url.URL{Scheme: "http", Host: "localhost:7777", Path: "/kite"}
go backendKite.Run()
<-backendKite.ServerReadyNotify()
// now search for a proxy from kontrol
color.Green("BackendKite is searching proxy from kontrol")
kites, err := backendKite.GetKites(&protocol.KontrolQuery{
Username: "******",
Environment: config.DefaultConfig.Environment,
Name: Name,
})
if err != nil {
t.Fatal(err)
}
proxyKite := kites[0]
err = proxyKite.Dial()
if err != nil {
t.Fatal(err)
}
// backendKite is connected to proxy, now let us register to proxy and get
// a proxy url. We send our url to proxy, it needs it in order to proxy us
color.Green("Backendkite found proxy, now registering to it")
result, err := proxyKite.TellWithTimeout("register", 4*time.Second, kiteUrl.String())
if err != nil {
t.Fatal(err)
}
proxyURL := result.MustString()
if !strings.Contains(proxyURL, "/proxy") {
t.Fatalf("Invalid proxy URL: %s", proxyURL)
}
registerURL, err := url.Parse(proxyURL)
if err != nil {
t.Fatal(err)
}
// register ourself to kontrol with this proxyUrl
color.Green("BackendKite is registering to Kontrol with the result from proxy")
go backendKite.RegisterForever(registerURL)
<-backendKite.KontrolReadyNotify()
// now another completely foreign kite and will search for our backend
// kite, connect to it and execute the "foo" method
color.Green("Foreign kite started")
foreignKite := kite.New("foreignKite", "1.0.0")
foreignKite.Config = conf.Copy()
color.Green("Querying backendKite now")
backendKites, err := foreignKite.GetKites(&protocol.KontrolQuery{
Username: "******",
Environment: config.DefaultConfig.Environment,
Name: "backendKite",
})
remoteBackendKite := backendKites[0]
color.Green("Dialing BackendKite")
err = remoteBackendKite.Dial()
if err != nil {
t.Fatal(err)
}
// foreignKite is connected to backendKite via proxy kite, fire our call...
color.Green("Calling BackendKite's foo method")
result, err = remoteBackendKite.TellWithTimeout("foo", 4*time.Second)
if err != nil {
t.Fatal(err)
}
s := result.MustString()
if s != "bar" {
t.Fatalf("Wrong reply: %s", s)
}
}