<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"path"
	"reflect"
	"testing"

	"golang.org/x/crypto/ssh"
)

type mockPublicKey struct {
	signature keySignature
}

func (publicKey mockPublicKey) Type() string {
	return publicKey.signature.String()
}

func (publicKey mockPublicKey) Marshal() []byte {
	return []byte(publicKey.signature.String())
}

func (publicKey mockPublicKey) Verify(data []byte, sig *ssh.Signature) error {
	return nil
}

type mockFile struct {
	closed bool
}

func (file *mockFile) Write(p []byte) (n int, err error) {
	return 0, errors.New("")
}

func (file *mockFile) Close() error {
	if file.closed {
		return errors.New("")
	}
	file.closed = true
	return nil
}

func verifyConfig(t *testing.T, cfg *config, expected *config) {
	if !reflect.DeepEqual(cfg.Server, expected.Server) {
		t.Errorf("Server=%v, want %v", cfg.Server, expected.Server)
	}
	if !reflect.DeepEqual(cfg.Logging, expected.Logging) {
		t.Errorf("Logging=%v, want %v", cfg.Logging, expected.Logging)
	}
	if !reflect.DeepEqual(cfg.Auth, expected.Auth) {
		t.Errorf("Auth=%v, want %v", cfg.Auth, expected.Auth)
	}
	if !reflect.DeepEqual(cfg.SSHProto, expected.SSHProto) {
		t.Errorf("SSHProto=%v, want %v", cfg.SSHProto, expected.SSHProto)
	}

	if cfg.sshConfig.RekeyThreshold != expected.SSHProto.RekeyThreshold {
		t.Errorf("sshConfig.RekeyThreshold=%v, want %v", cfg.sshConfig.RekeyThreshold, expected.SSHProto.RekeyThreshold)
	}
	if !reflect.DeepEqual(cfg.sshConfig.KeyExchanges, expected.SSHProto.KeyExchanges) {
		t.Errorf("sshConfig.KeyExchanges=%v, want %v", cfg.sshConfig.KeyExchanges, expected.SSHProto.KeyExchanges)
	}
	if !reflect.DeepEqual(cfg.sshConfig.Ciphers, expected.SSHProto.Ciphers) {
		t.Errorf("sshConfig.Ciphers=%v, want %v", cfg.sshConfig.Ciphers, expected.SSHProto.Ciphers)
	}
	if !reflect.DeepEqual(cfg.sshConfig.MACs, expected.SSHProto.MACs) {
		t.Errorf("sshConfig.MACs=%v, want %v", cfg.sshConfig.MACs, expected.SSHProto.MACs)
	}
	if cfg.sshConfig.NoClientAuth != expected.Auth.NoAuth {
		t.Errorf("sshConfig.NoClientAuth=%v, want %v", cfg.sshConfig.NoClientAuth, expected.Auth.NoAuth)
	}
	if cfg.sshConfig.MaxAuthTries != expected.Auth.MaxTries {
		t.Errorf("sshConfig.MaxAuthTries=%v, want %v", cfg.sshConfig.MaxAuthTries, expected.Auth.MaxTries)
	}
	if (cfg.sshConfig.PasswordCallback != nil) != expected.Auth.PasswordAuth.Enabled {
		t.Errorf("sshConfig.PasswordCallback=%v, want %v", cfg.sshConfig.PasswordCallback != nil, expected.Auth.PasswordAuth.Enabled)
	}
	if (cfg.sshConfig.PublicKeyCallback != nil) != expected.Auth.PublicKeyAuth.Enabled {
		t.Errorf("sshConfig.PasswordCallback=%v, want %v", cfg.sshConfig.PublicKeyCallback != nil, expected.Auth.PublicKeyAuth.Enabled)
	}
	if (cfg.sshConfig.KeyboardInteractiveCallback != nil) != expected.Auth.KeyboardInteractiveAuth.Enabled {
		t.Errorf("sshConfig.KeyboardInteractiveCallback=%v, want %v", cfg.sshConfig.KeyboardInteractiveCallback != nil, expected.Auth.KeyboardInteractiveAuth.Enabled)
	}
	if cfg.sshConfig.AuthLogCallback == nil {
		t.Errorf("sshConfig.AuthLogCallback=nil, want a callback")
	}
	if cfg.sshConfig.ServerVersion != expected.SSHProto.Version {
		t.Errorf("sshConfig.ServerVersion=%v, want %v", cfg.sshConfig.ServerVersion, expected.SSHProto.Version)
	}
	if (cfg.sshConfig.BannerCallback != nil) != (expected.SSHProto.Banner != "") {
		t.Errorf("sshConfig.BannerCallback=%v, want %v", cfg.sshConfig.BannerCallback != nil, expected.SSHProto.Banner != "")
	}
	if cfg.sshConfig.GSSAPIWithMICConfig != nil {
		t.Errorf("sshConfig.GSSAPIWithMICConfig=%v, want nil", cfg.sshConfig.GSSAPIWithMICConfig)
	}
	if len(cfg.parsedHostKeys) != len(expected.Server.HostKeys) {
		t.Errorf("len(parsedHostKeys)=%v, want %v", len(cfg.parsedHostKeys), len(expected.Server.HostKeys))
	}

	if expected.Logging.File == "" {
		if cfg.logFileHandle != nil {
			t.Errorf("logFileHandle=%v, want nil", cfg.logFileHandle)
		}
	} else {
		if cfg.logFileHandle == nil {
			t.Errorf("logFileHandle=nil, want a file")
		}
	}
}

func verifyDefaultKeys(t *testing.T, dataDir string) {
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		t.Fatalf("Faield to list directory: %v", err)
	}
	expectedKeys := map[string]string{
		"host_rsa_key":     "ssh-rsa",
		"host_ecdsa_key":   "ecdsa-sha2-nistp256",
		"host_ed25519_key": "ssh-ed25519",
	}
	keys := map[string]string{}
	for _, file := range files {
		keyBytes, err := ioutil.ReadFile(path.Join(dataDir, file.Name()))
		if err != nil {
			t.Fatalf("Failed to read key: %v", err)
		}
		signer, err := ssh.ParsePrivateKey(keyBytes)
		if err != nil {
			t.Fatalf("Failed to parse private key: %v", err)
		}
		keys[file.Name()] = signer.PublicKey().Type()
	}
	if !reflect.DeepEqual(keys, expectedKeys) {
		t.Errorf("keys=%v, want %v", keys, expectedKeys)
	}
}

func TestDefaultConfig(t *testing.T) {
	dataDir := t.TempDir()
	cfg, err := getConfig("", dataDir)
	if err != nil {
		t.Fatalf("Failed to get config: %v", err)
	}
	expectedConfig := &config{}
	expectedConfig.Server.ListenAddress = "127.0.0.1:2022"
	expectedConfig.Server.HostKeys = []string{
		path.Join(dataDir, "host_rsa_key"),
		path.Join(dataDir, "host_ecdsa_key"),
		path.Join(dataDir, "host_ed25519_key"),
	}
	expectedConfig.Server.TCPIPServices = map[uint32]string{
		25:   "SMTP",
		80:   "HTTP",
		110:  "POP3",
		587:  "SMTP",
		8080: "HTTP",
	}
	expectedConfig.Logging.Timestamps = true
	expectedConfig.Auth.PasswordAuth.Enabled = true
	expectedConfig.Auth.PasswordAuth.Accepted = true
	expectedConfig.Auth.PublicKeyAuth.Enabled = true
	expectedConfig.SSHProto.Version = "SSH-2.0-sshesame"
	expectedConfig.SSHProto.Banner = "This is an SSH honeypot. Everything is logged and monitored."
	verifyConfig(t, cfg, expectedConfig)
	verifyDefaultKeys(t, dataDir)
}

func TestUserConfigDefaultKeys(t *testing.T) {
	logFile := path.Join(t.TempDir(), "test.log")
	cfgString := fmt.Sprintf(`
server:
  listen_address: 0.0.0.0:22
  tcpip_services: {}
logging:
  file: %v
  json: true
  timestamps: false
  metrics_address: 0.0.0.0:2112
auth:
  max_tries: 234
  no_auth: true
  password_auth:
    enabled: false
    accepted: false
  public_key_auth:
    enabled: false
    accepted: true
  keyboard_interactive_auth:
    enabled: true
    accepted: true
    instruction: instruction
    questions:
    - text: q1
      echo: true
    - text: q2
      echo: false
ssh_proto:
  version: SSH-2.0-test
  banner:
  rekey_threshold: 123
  key_exchanges: [kex]
  ciphers: [cipher]
  macs: [mac]
`, logFile)
	dataDir := t.TempDir()
	cfg, err := getConfig(cfgString, dataDir)
	if err != nil {
		t.Fatalf("Failed to get config: %v", err)
	}
	if cfg.logFileHandle != nil {
		cfg.logFileHandle.Close()
	}
	expectedConfig := &config{}
	expectedConfig.Server.ListenAddress = "0.0.0.0:22"
	expectedConfig.Server.HostKeys = []string{
		path.Join(dataDir, "host_rsa_key"),
		path.Join(dataDir, "host_ecdsa_key"),
		path.Join(dataDir, "host_ed25519_key"),
	}
	expectedConfig.Server.TCPIPServices = map[uint32]string{}
	expectedConfig.Logging.File = logFile
	expectedConfig.Logging.JSON = true
	expectedConfig.Logging.Timestamps = false
	expectedConfig.Logging.MetricsAddress = "0.0.0.0:2112"
	expectedConfig.Auth.MaxTries = 234
	expectedConfig.Auth.NoAuth = true
	expectedConfig.Auth.PublicKeyAuth.Accepted = true
	expectedConfig.Auth.KeyboardInteractiveAuth.Enabled = true
	expectedConfig.Auth.KeyboardInteractiveAuth.Accepted = true
	expectedConfig.Auth.KeyboardInteractiveAuth.Instruction = "instruction"
	expectedConfig.Auth.KeyboardInteractiveAuth.Questions = []keyboardInteractiveAuthQuestion{
		{"q1", true},
		{"q2", false},
	}
	expectedConfig.SSHProto.Version = "SSH-2.0-test"
	expectedConfig.SSHProto.RekeyThreshold = 123
	expectedConfig.SSHProto.KeyExchanges = []string{"kex"}
	expectedConfig.SSHProto.Ciphers = []string{"cipher"}
	expectedConfig.SSHProto.MACs = []string{"mac"}
	verifyConfig(t, cfg, expectedConfig)
	verifyDefaultKeys(t, dataDir)
}

func TestUserConfigCustomKeysAndServices(t *testing.T) {
	keyFile, err := generateKey(t.TempDir(), ecdsa_key)
	cfgString := fmt.Sprintf(`
server:
  host_keys: [%v]
  tcpip_services:
    8080: HTTP
`, keyFile)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	dataDir := t.TempDir()
	cfg, err := getConfig(cfgString, dataDir)
	if err != nil {
		t.Fatalf("Failed to get config: %v", err)
	}
	expectedConfig := &config{}
	expectedConfig.Server.ListenAddress = "127.0.0.1:2022"
	expectedConfig.Server.HostKeys = []string{keyFile}
	expectedConfig.Server.TCPIPServices = map[uint32]string{
		8080: "HTTP",
	}
	expectedConfig.Logging.Timestamps = true
	expectedConfig.Auth.PasswordAuth.Enabled = true
	expectedConfig.Auth.PasswordAuth.Accepted = true
	expectedConfig.Auth.PublicKeyAuth.Enabled = true
	expectedConfig.SSHProto.Version = "SSH-2.0-sshesame"
	expectedConfig.SSHProto.Banner = "This is an SSH honeypot. Everything is logged and monitored."
	verifyConfig(t, cfg, expectedConfig)
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		t.Fatalf("Failed to read directory: %v", err)
	}
	if len(files) != 0 {
		t.Errorf("files=%v, want []", files)
	}
}

func TestSetupLoggingOldHandleClosed(t *testing.T) {
	file := &mockFile{}
	cfg := &config{logFileHandle: file}
	if err := cfg.setupLogging(); err != nil {
		t.Fatalf("Failed to set up logging: %v", err)
	}
	if !file.closed {
		t.Errorf("file.closed=false, want true")
	}
}

func TestExistingKey(t *testing.T) {
	dataDir := path.Join(t.TempDir(), "keys")
	oldKeyFile, err := generateKey(dataDir, ed25519_key)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	oldKey, err := ioutil.ReadFile(oldKeyFile)
	if err != nil {
		t.Fatalf("Failed to read key: %v", err)
	}
	newKeyFile, err := generateKey(dataDir, ed25519_key)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	newKey, err := ioutil.ReadFile(newKeyFile)
	if err != nil {
		t.Fatalf("Failed to read key: %v", err)
	}
	if !reflect.DeepEqual(oldKey, newKey) {
		t.Errorf("oldKey!=newKey")
	}
}

func TestDefaultConfigFile(t *testing.T) {
	configBytes, err := ioutil.ReadFile("sshesame.yaml")
	if err != nil {
		t.Fatalf("Failed to read config file: %v", err)
	}
	dataDir := t.TempDir()
	cfg, err := getConfig(string(configBytes), dataDir)
	if err != nil {
		t.Fatalf("Failed to get config: %v", err)
	}

	// The sample config has example keyboard interactive auth options set.
	// Since the auth method itself is disabled, this doesn't make a difference.
	// Unset them so they don't affect the comparison.
	cfg.Auth.KeyboardInteractiveAuth.Instruction = ""
	cfg.Auth.KeyboardInteractiveAuth.Questions = nil

	defaultCfg, err := getConfig("", dataDir)
	if err != nil {
		t.Fatalf("Failed to get default config: %v", err)
	}
	verifyConfig(t, cfg, defaultCfg)
}
<commit_msg>TestDefaultConfigFile: don't expand the parsed config with default options<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"path"
	"reflect"
	"testing"

	"golang.org/x/crypto/ssh"
	"gopkg.in/yaml.v2"
)

type mockPublicKey struct {
	signature keySignature
}

func (publicKey mockPublicKey) Type() string {
	return publicKey.signature.String()
}

func (publicKey mockPublicKey) Marshal() []byte {
	return []byte(publicKey.signature.String())
}

func (publicKey mockPublicKey) Verify(data []byte, sig *ssh.Signature) error {
	return nil
}

type mockFile struct {
	closed bool
}

func (file *mockFile) Write(p []byte) (n int, err error) {
	return 0, errors.New("")
}

func (file *mockFile) Close() error {
	if file.closed {
		return errors.New("")
	}
	file.closed = true
	return nil
}

func verifyConfig(t *testing.T, cfg *config, expected *config) {
	if !reflect.DeepEqual(cfg.Server, expected.Server) {
		t.Errorf("Server=%v, want %v", cfg.Server, expected.Server)
	}
	if !reflect.DeepEqual(cfg.Logging, expected.Logging) {
		t.Errorf("Logging=%v, want %v", cfg.Logging, expected.Logging)
	}
	if !reflect.DeepEqual(cfg.Auth, expected.Auth) {
		t.Errorf("Auth=%v, want %v", cfg.Auth, expected.Auth)
	}
	if !reflect.DeepEqual(cfg.SSHProto, expected.SSHProto) {
		t.Errorf("SSHProto=%v, want %v", cfg.SSHProto, expected.SSHProto)
	}

	if cfg.sshConfig.RekeyThreshold != expected.SSHProto.RekeyThreshold {
		t.Errorf("sshConfig.RekeyThreshold=%v, want %v", cfg.sshConfig.RekeyThreshold, expected.SSHProto.RekeyThreshold)
	}
	if !reflect.DeepEqual(cfg.sshConfig.KeyExchanges, expected.SSHProto.KeyExchanges) {
		t.Errorf("sshConfig.KeyExchanges=%v, want %v", cfg.sshConfig.KeyExchanges, expected.SSHProto.KeyExchanges)
	}
	if !reflect.DeepEqual(cfg.sshConfig.Ciphers, expected.SSHProto.Ciphers) {
		t.Errorf("sshConfig.Ciphers=%v, want %v", cfg.sshConfig.Ciphers, expected.SSHProto.Ciphers)
	}
	if !reflect.DeepEqual(cfg.sshConfig.MACs, expected.SSHProto.MACs) {
		t.Errorf("sshConfig.MACs=%v, want %v", cfg.sshConfig.MACs, expected.SSHProto.MACs)
	}
	if cfg.sshConfig.NoClientAuth != expected.Auth.NoAuth {
		t.Errorf("sshConfig.NoClientAuth=%v, want %v", cfg.sshConfig.NoClientAuth, expected.Auth.NoAuth)
	}
	if cfg.sshConfig.MaxAuthTries != expected.Auth.MaxTries {
		t.Errorf("sshConfig.MaxAuthTries=%v, want %v", cfg.sshConfig.MaxAuthTries, expected.Auth.MaxTries)
	}
	if (cfg.sshConfig.PasswordCallback != nil) != expected.Auth.PasswordAuth.Enabled {
		t.Errorf("sshConfig.PasswordCallback=%v, want %v", cfg.sshConfig.PasswordCallback != nil, expected.Auth.PasswordAuth.Enabled)
	}
	if (cfg.sshConfig.PublicKeyCallback != nil) != expected.Auth.PublicKeyAuth.Enabled {
		t.Errorf("sshConfig.PasswordCallback=%v, want %v", cfg.sshConfig.PublicKeyCallback != nil, expected.Auth.PublicKeyAuth.Enabled)
	}
	if (cfg.sshConfig.KeyboardInteractiveCallback != nil) != expected.Auth.KeyboardInteractiveAuth.Enabled {
		t.Errorf("sshConfig.KeyboardInteractiveCallback=%v, want %v", cfg.sshConfig.KeyboardInteractiveCallback != nil, expected.Auth.KeyboardInteractiveAuth.Enabled)
	}
	if cfg.sshConfig.AuthLogCallback == nil {
		t.Errorf("sshConfig.AuthLogCallback=nil, want a callback")
	}
	if cfg.sshConfig.ServerVersion != expected.SSHProto.Version {
		t.Errorf("sshConfig.ServerVersion=%v, want %v", cfg.sshConfig.ServerVersion, expected.SSHProto.Version)
	}
	if (cfg.sshConfig.BannerCallback != nil) != (expected.SSHProto.Banner != "") {
		t.Errorf("sshConfig.BannerCallback=%v, want %v", cfg.sshConfig.BannerCallback != nil, expected.SSHProto.Banner != "")
	}
	if cfg.sshConfig.GSSAPIWithMICConfig != nil {
		t.Errorf("sshConfig.GSSAPIWithMICConfig=%v, want nil", cfg.sshConfig.GSSAPIWithMICConfig)
	}
	if len(cfg.parsedHostKeys) != len(expected.Server.HostKeys) {
		t.Errorf("len(parsedHostKeys)=%v, want %v", len(cfg.parsedHostKeys), len(expected.Server.HostKeys))
	}

	if expected.Logging.File == "" {
		if cfg.logFileHandle != nil {
			t.Errorf("logFileHandle=%v, want nil", cfg.logFileHandle)
		}
	} else {
		if cfg.logFileHandle == nil {
			t.Errorf("logFileHandle=nil, want a file")
		}
	}
}

func verifyDefaultKeys(t *testing.T, dataDir string) {
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		t.Fatalf("Faield to list directory: %v", err)
	}
	expectedKeys := map[string]string{
		"host_rsa_key":     "ssh-rsa",
		"host_ecdsa_key":   "ecdsa-sha2-nistp256",
		"host_ed25519_key": "ssh-ed25519",
	}
	keys := map[string]string{}
	for _, file := range files {
		keyBytes, err := ioutil.ReadFile(path.Join(dataDir, file.Name()))
		if err != nil {
			t.Fatalf("Failed to read key: %v", err)
		}
		signer, err := ssh.ParsePrivateKey(keyBytes)
		if err != nil {
			t.Fatalf("Failed to parse private key: %v", err)
		}
		keys[file.Name()] = signer.PublicKey().Type()
	}
	if !reflect.DeepEqual(keys, expectedKeys) {
		t.Errorf("keys=%v, want %v", keys, expectedKeys)
	}
}

func TestDefaultConfig(t *testing.T) {
	dataDir := t.TempDir()
	cfg, err := getConfig("", dataDir)
	if err != nil {
		t.Fatalf("Failed to get config: %v", err)
	}
	expectedConfig := &config{}
	expectedConfig.Server.ListenAddress = "127.0.0.1:2022"
	expectedConfig.Server.HostKeys = []string{
		path.Join(dataDir, "host_rsa_key"),
		path.Join(dataDir, "host_ecdsa_key"),
		path.Join(dataDir, "host_ed25519_key"),
	}
	expectedConfig.Server.TCPIPServices = map[uint32]string{
		25:   "SMTP",
		80:   "HTTP",
		110:  "POP3",
		587:  "SMTP",
		8080: "HTTP",
	}
	expectedConfig.Logging.Timestamps = true
	expectedConfig.Auth.PasswordAuth.Enabled = true
	expectedConfig.Auth.PasswordAuth.Accepted = true
	expectedConfig.Auth.PublicKeyAuth.Enabled = true
	expectedConfig.SSHProto.Version = "SSH-2.0-sshesame"
	expectedConfig.SSHProto.Banner = "This is an SSH honeypot. Everything is logged and monitored."
	verifyConfig(t, cfg, expectedConfig)
	verifyDefaultKeys(t, dataDir)
}

func TestUserConfigDefaultKeys(t *testing.T) {
	logFile := path.Join(t.TempDir(), "test.log")
	cfgString := fmt.Sprintf(`
server:
  listen_address: 0.0.0.0:22
  tcpip_services: {}
logging:
  file: %v
  json: true
  timestamps: false
  metrics_address: 0.0.0.0:2112
auth:
  max_tries: 234
  no_auth: true
  password_auth:
    enabled: false
    accepted: false
  public_key_auth:
    enabled: false
    accepted: true
  keyboard_interactive_auth:
    enabled: true
    accepted: true
    instruction: instruction
    questions:
    - text: q1
      echo: true
    - text: q2
      echo: false
ssh_proto:
  version: SSH-2.0-test
  banner:
  rekey_threshold: 123
  key_exchanges: [kex]
  ciphers: [cipher]
  macs: [mac]
`, logFile)
	dataDir := t.TempDir()
	cfg, err := getConfig(cfgString, dataDir)
	if err != nil {
		t.Fatalf("Failed to get config: %v", err)
	}
	if cfg.logFileHandle != nil {
		cfg.logFileHandle.Close()
	}
	expectedConfig := &config{}
	expectedConfig.Server.ListenAddress = "0.0.0.0:22"
	expectedConfig.Server.HostKeys = []string{
		path.Join(dataDir, "host_rsa_key"),
		path.Join(dataDir, "host_ecdsa_key"),
		path.Join(dataDir, "host_ed25519_key"),
	}
	expectedConfig.Server.TCPIPServices = map[uint32]string{}
	expectedConfig.Logging.File = logFile
	expectedConfig.Logging.JSON = true
	expectedConfig.Logging.Timestamps = false
	expectedConfig.Logging.MetricsAddress = "0.0.0.0:2112"
	expectedConfig.Auth.MaxTries = 234
	expectedConfig.Auth.NoAuth = true
	expectedConfig.Auth.PublicKeyAuth.Accepted = true
	expectedConfig.Auth.KeyboardInteractiveAuth.Enabled = true
	expectedConfig.Auth.KeyboardInteractiveAuth.Accepted = true
	expectedConfig.Auth.KeyboardInteractiveAuth.Instruction = "instruction"
	expectedConfig.Auth.KeyboardInteractiveAuth.Questions = []keyboardInteractiveAuthQuestion{
		{"q1", true},
		{"q2", false},
	}
	expectedConfig.SSHProto.Version = "SSH-2.0-test"
	expectedConfig.SSHProto.RekeyThreshold = 123
	expectedConfig.SSHProto.KeyExchanges = []string{"kex"}
	expectedConfig.SSHProto.Ciphers = []string{"cipher"}
	expectedConfig.SSHProto.MACs = []string{"mac"}
	verifyConfig(t, cfg, expectedConfig)
	verifyDefaultKeys(t, dataDir)
}

func TestUserConfigCustomKeysAndServices(t *testing.T) {
	keyFile, err := generateKey(t.TempDir(), ecdsa_key)
	cfgString := fmt.Sprintf(`
server:
  host_keys: [%v]
  tcpip_services:
    8080: HTTP
`, keyFile)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	dataDir := t.TempDir()
	cfg, err := getConfig(cfgString, dataDir)
	if err != nil {
		t.Fatalf("Failed to get config: %v", err)
	}
	expectedConfig := &config{}
	expectedConfig.Server.ListenAddress = "127.0.0.1:2022"
	expectedConfig.Server.HostKeys = []string{keyFile}
	expectedConfig.Server.TCPIPServices = map[uint32]string{
		8080: "HTTP",
	}
	expectedConfig.Logging.Timestamps = true
	expectedConfig.Auth.PasswordAuth.Enabled = true
	expectedConfig.Auth.PasswordAuth.Accepted = true
	expectedConfig.Auth.PublicKeyAuth.Enabled = true
	expectedConfig.SSHProto.Version = "SSH-2.0-sshesame"
	expectedConfig.SSHProto.Banner = "This is an SSH honeypot. Everything is logged and monitored."
	verifyConfig(t, cfg, expectedConfig)
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		t.Fatalf("Failed to read directory: %v", err)
	}
	if len(files) != 0 {
		t.Errorf("files=%v, want []", files)
	}
}

func TestSetupLoggingOldHandleClosed(t *testing.T) {
	file := &mockFile{}
	cfg := &config{logFileHandle: file}
	if err := cfg.setupLogging(); err != nil {
		t.Fatalf("Failed to set up logging: %v", err)
	}
	if !file.closed {
		t.Errorf("file.closed=false, want true")
	}
}

func TestExistingKey(t *testing.T) {
	dataDir := path.Join(t.TempDir(), "keys")
	oldKeyFile, err := generateKey(dataDir, ed25519_key)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	oldKey, err := ioutil.ReadFile(oldKeyFile)
	if err != nil {
		t.Fatalf("Failed to read key: %v", err)
	}
	newKeyFile, err := generateKey(dataDir, ed25519_key)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	newKey, err := ioutil.ReadFile(newKeyFile)
	if err != nil {
		t.Fatalf("Failed to read key: %v", err)
	}
	if !reflect.DeepEqual(oldKey, newKey) {
		t.Errorf("oldKey!=newKey")
	}
}

func TestDefaultConfigFile(t *testing.T) {
	configBytes, err := ioutil.ReadFile("sshesame.yaml")
	if err != nil {
		t.Fatalf("Failed to read config file: %v", err)
	}
	cfg := &config{}
	if err := yaml.UnmarshalStrict(configBytes, cfg); err != nil {
		t.Fatalf("Failed to unmarshal config: %v", err)
	}
	dataDir := t.TempDir()
	if err := cfg.setDefaultHostKeys(dataDir, []keySignature{rsa_key, ecdsa_key, ed25519_key}); err != nil {
		t.Fatalf("Failed to set default host keys: %v", err)
	}
	if err := cfg.setupSSHConfig(); err != nil {
		t.Fatalf("Failed to setup SSH config: %v", err)
	}

	// The sample config has example keyboard interactive auth options set.
	// Since the auth method itself is disabled, this doesn't make a difference.
	// Unset them so they don't affect the comparison.
	cfg.Auth.KeyboardInteractiveAuth.Instruction = ""
	cfg.Auth.KeyboardInteractiveAuth.Questions = nil

	defaultCfg, err := getConfig("", dataDir)
	if err != nil {
		t.Fatalf("Failed to get default config: %v", err)
	}
	verifyConfig(t, cfg, defaultCfg)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/consul-template/test"
	"github.com/hashicorp/consul-template/util"
)

// Test that an empty config does nothing
func TestMerge_emptyConfig(t *testing.T) {
	consul := "consul.io:8500"
	config := &Config{Consul: consul}
	config.Merge(&Config{})

	if config.Consul != consul {
		t.Fatalf("expected %q to equal %q", config.Consul, consul)
	}
}

// Test that simple values are merged
func TestMerge_simpleConfig(t *testing.T) {
	config, newConsul := &Config{Consul: "consul.io:8500"}, "packer.io:7300"
	config.Merge(&Config{Consul: newConsul})

	if config.Consul != newConsul {
		t.Fatalf("expected %q to equal %q", config.Consul, newConsul)
	}
}

// Test that complex values are merged
func TestMerge_complexConfig(t *testing.T) {
	templates := []*ConfigTemplate{
		&ConfigTemplate{
			Source:      "a",
			Destination: "b",
		},
		&ConfigTemplate{
			Source:      "c",
			Destination: "d",
			Command:     "e",
		},
		&ConfigTemplate{
			Source:      "f",
			Destination: "g",
			Command:     "h",
		},
		&ConfigTemplate{
			Source:      "i",
			Destination: "j",
		},
	}

	config := &Config{ConfigTemplates: templates[:2]}
	config.Merge(&Config{ConfigTemplates: templates[2:]})

	if !reflect.DeepEqual(config.ConfigTemplates, templates) {
		t.Fatalf("expected %q to equal %q", config.ConfigTemplates, templates)
	}
}

// Test that file read errors are propagated up
func TestParseConfig_readFileError(t *testing.T) {
	_, err := ParseConfig(path.Join(os.TempDir(), "config.json"))
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "no such file or directory"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that parser errors are propagated up
func TestParseConfig_parseFileError(t *testing.T) {
	configFile := test.CreateTempfile([]byte(`
    invalid file in here
  `), t)
	defer test.DeleteTempfile(configFile, t)

	_, err := ParseConfig(configFile.Name())
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "syntax error"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that mapstructure errors are propagated up
func TestParseConfig_mapstructureError(t *testing.T) {
	configFile := test.CreateTempfile([]byte(`
    consul = true
  `), t)
	defer test.DeleteTempfile(configFile, t)

	_, err := ParseConfig(configFile.Name())
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "nconvertible type 'bool'"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that the config is parsed correctly
func TestParseConfig_correctValues(t *testing.T) {
	configFile := test.CreateTempfile([]byte(`
    consul = "nyc1.demo.consul.io"
    token = "abcd1234"
    wait = "5s:10s"

    template {
      source = "nginx.conf.ctmpl"
      destination  = "/etc/nginx/nginx.conf"
    }

    template {
      source = "redis.conf.ctmpl"
      destination  = "/etc/redis/redis.conf"
      command = "service redis restart"
    }
  `), t)
	defer test.DeleteTempfile(configFile, t)

	config, err := ParseConfig(configFile.Name())
	if err != nil {
		t.Fatal(err)
	}

	expected := &Config{
		Path:   configFile.Name(),
		Consul: "nyc1.demo.consul.io",
		Token:  "abcd1234",
		Wait: &util.Wait{
			Min: time.Second * 5,
			Max: time.Second * 10,
		},
		WaitRaw: "5s:10s",
		ConfigTemplates: []*ConfigTemplate{
			&ConfigTemplate{
				Source:      "nginx.conf.ctmpl",
				Destination: "/etc/nginx/nginx.conf",
			},
			&ConfigTemplate{
				Source:      "redis.conf.ctmpl",
				Destination: "/etc/redis/redis.conf",
				Command:     "service redis restart",
			},
		},
	}
	if !reflect.DeepEqual(config, expected) {
		t.Fatalf("expected %+v to be %+v", expected, config)
	}
}

// Test that ParseWait errors are propagated up
func TestParseConfig_parseWaitError(t *testing.T) {
	configFile := test.CreateTempfile([]byte(`
    wait = "not_valid:duration"
  `), t)
	defer test.DeleteTempfile(configFile, t)

	_, err := ParseConfig(configFile.Name())
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "invalid duration not_valid"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that an error is returned when the empty string is given
func TestParseConfigTemplate_emptyStringArgs(t *testing.T) {
	_, err := ParseConfigTemplate("")
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "cannot specify empty template declaration"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that an error is returned when a string with spaces is given
func TestParseConfigTemplate_stringWithSpacesArgs(t *testing.T) {
	_, err := ParseConfigTemplate("  ")
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "cannot specify empty template declaration"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that an error is returned when there are too many arguments
func TestParseConfigurationTemplate_tooManyArgs(t *testing.T) {
	_, err := ParseConfigTemplate("foo:bar:blitz:baz")
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "invalid template declaration format"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that a source value is correctly used
func TestParseConfigurationTemplate_source(t *testing.T) {
	source := "/tmp/config.ctmpl"
	template, err := ParseConfigTemplate(source)
	if err != nil {
		t.Fatal(err)
	}

	if template.Source != source {
		t.Errorf("expected %q to equal %q", template.Source, source)
	}
}

// Test that a destination wait value is correctly used
func TestParseConfigurationTemplate_destination(t *testing.T) {
	source, destination := "/tmp/config.ctmpl", "/tmp/out"
	template, err := ParseConfigTemplate(fmt.Sprintf("%s:%s", source, destination))
	if err != nil {
		t.Fatal(err)
	}

	if template.Source != source {
		t.Errorf("expected %q to equal %q", template.Source, source)
	}

	if template.Destination != destination {
		t.Errorf("expected %q to equal %q", template.Destination, destination)
	}
}

// Test that a command wait value is correctly used
func TestParseConfigurationTemplate_command(t *testing.T) {
	source, destination, command := "/tmp/config.ctmpl", "/tmp/out", "reboot"
	template, err := ParseConfigTemplate(fmt.Sprintf("%s:%s:%s", source, destination, command))
	if err != nil {
		t.Fatal(err)
	}

	if template.Source != source {
		t.Errorf("expected %q to equal %q", template.Source, source)
	}

	if template.Destination != destination {
		t.Errorf("expected %q to equal %q", template.Destination, destination)
	}

	if template.Command != command {
		t.Errorf("expected %q to equal %q", template.Command, command)
	}
}
<commit_msg>Update merge test to be clearer<commit_after>package main

import (
	"fmt"
	"os"
	"path"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/consul-template/test"
	"github.com/hashicorp/consul-template/util"
)

// Test that an empty config does nothing
func TestMerge_emptyConfig(t *testing.T) {
	consul := "consul.io:8500"
	config := &Config{Consul: consul}
	config.Merge(&Config{})

	if config.Consul != consul {
		t.Fatalf("expected %q to equal %q", config.Consul, consul)
	}
}

// Test that simple values are merged
func TestMerge_simpleConfig(t *testing.T) {
	config, newConsul := &Config{Consul: "consul.io:8500"}, "packer.io:7300"
	config.Merge(&Config{Consul: newConsul})

	if config.Consul != newConsul {
		t.Fatalf("expected %q to equal %q", config.Consul, newConsul)
	}
}

// Test that complex values are merged, and that ConfigTemplates are additive
func TestMerge_complexConfig(t *testing.T) {
	templates := []*ConfigTemplate{
		&ConfigTemplate{
			Source:      "a",
			Destination: "b",
		},
		&ConfigTemplate{
			Source:      "c",
			Destination: "d",
			Command:     "e",
		},
		&ConfigTemplate{
			Source:      "f",
			Destination: "g",
			Command:     "h",
		},
		&ConfigTemplate{
			Source:      "i",
			Destination: "j",
		},
	}

	config := &Config{ConfigTemplates: templates[:2]}
	otherConfig := &Config{ConfigTemplates: templates[2:]}
	config.Merge(otherConfig)

	if !reflect.DeepEqual(config.ConfigTemplates, templates) {
		t.Fatalf("expected %q to equal %q", config.ConfigTemplates, templates)
	}
}

// Test that file read errors are propagated up
func TestParseConfig_readFileError(t *testing.T) {
	_, err := ParseConfig(path.Join(os.TempDir(), "config.json"))
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "no such file or directory"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that parser errors are propagated up
func TestParseConfig_parseFileError(t *testing.T) {
	configFile := test.CreateTempfile([]byte(`
    invalid file in here
  `), t)
	defer test.DeleteTempfile(configFile, t)

	_, err := ParseConfig(configFile.Name())
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "syntax error"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that mapstructure errors are propagated up
func TestParseConfig_mapstructureError(t *testing.T) {
	configFile := test.CreateTempfile([]byte(`
    consul = true
  `), t)
	defer test.DeleteTempfile(configFile, t)

	_, err := ParseConfig(configFile.Name())
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "nconvertible type 'bool'"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that the config is parsed correctly
func TestParseConfig_correctValues(t *testing.T) {
	configFile := test.CreateTempfile([]byte(`
    consul = "nyc1.demo.consul.io"
    token = "abcd1234"
    wait = "5s:10s"

    template {
      source = "nginx.conf.ctmpl"
      destination  = "/etc/nginx/nginx.conf"
    }

    template {
      source = "redis.conf.ctmpl"
      destination  = "/etc/redis/redis.conf"
      command = "service redis restart"
    }
  `), t)
	defer test.DeleteTempfile(configFile, t)

	config, err := ParseConfig(configFile.Name())
	if err != nil {
		t.Fatal(err)
	}

	expected := &Config{
		Path:   configFile.Name(),
		Consul: "nyc1.demo.consul.io",
		Token:  "abcd1234",
		Wait: &util.Wait{
			Min: time.Second * 5,
			Max: time.Second * 10,
		},
		WaitRaw: "5s:10s",
		ConfigTemplates: []*ConfigTemplate{
			&ConfigTemplate{
				Source:      "nginx.conf.ctmpl",
				Destination: "/etc/nginx/nginx.conf",
			},
			&ConfigTemplate{
				Source:      "redis.conf.ctmpl",
				Destination: "/etc/redis/redis.conf",
				Command:     "service redis restart",
			},
		},
	}
	if !reflect.DeepEqual(config, expected) {
		t.Fatalf("expected %+v to be %+v", expected, config)
	}
}

// Test that ParseWait errors are propagated up
func TestParseConfig_parseWaitError(t *testing.T) {
	configFile := test.CreateTempfile([]byte(`
    wait = "not_valid:duration"
  `), t)
	defer test.DeleteTempfile(configFile, t)

	_, err := ParseConfig(configFile.Name())
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "invalid duration not_valid"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that an error is returned when the empty string is given
func TestParseConfigTemplate_emptyStringArgs(t *testing.T) {
	_, err := ParseConfigTemplate("")
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "cannot specify empty template declaration"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that an error is returned when a string with spaces is given
func TestParseConfigTemplate_stringWithSpacesArgs(t *testing.T) {
	_, err := ParseConfigTemplate("  ")
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "cannot specify empty template declaration"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that an error is returned when there are too many arguments
func TestParseConfigurationTemplate_tooManyArgs(t *testing.T) {
	_, err := ParseConfigTemplate("foo:bar:blitz:baz")
	if err == nil {
		t.Fatal("expected error, but nothing was returned")
	}

	expectedErr := "invalid template declaration format"
	if !strings.Contains(err.Error(), expectedErr) {
		t.Fatalf("expected error %q to contain %q", err.Error(), expectedErr)
	}
}

// Test that a source value is correctly used
func TestParseConfigurationTemplate_source(t *testing.T) {
	source := "/tmp/config.ctmpl"
	template, err := ParseConfigTemplate(source)
	if err != nil {
		t.Fatal(err)
	}

	if template.Source != source {
		t.Errorf("expected %q to equal %q", template.Source, source)
	}
}

// Test that a destination wait value is correctly used
func TestParseConfigurationTemplate_destination(t *testing.T) {
	source, destination := "/tmp/config.ctmpl", "/tmp/out"
	template, err := ParseConfigTemplate(fmt.Sprintf("%s:%s", source, destination))
	if err != nil {
		t.Fatal(err)
	}

	if template.Source != source {
		t.Errorf("expected %q to equal %q", template.Source, source)
	}

	if template.Destination != destination {
		t.Errorf("expected %q to equal %q", template.Destination, destination)
	}
}

// Test that a command wait value is correctly used
func TestParseConfigurationTemplate_command(t *testing.T) {
	source, destination, command := "/tmp/config.ctmpl", "/tmp/out", "reboot"
	template, err := ParseConfigTemplate(fmt.Sprintf("%s:%s:%s", source, destination, command))
	if err != nil {
		t.Fatal(err)
	}

	if template.Source != source {
		t.Errorf("expected %q to equal %q", template.Source, source)
	}

	if template.Destination != destination {
		t.Errorf("expected %q to equal %q", template.Destination, destination)
	}

	if template.Command != command {
		t.Errorf("expected %q to equal %q", template.Command, command)
	}
}
<|endoftext|>


<commit_before>package consul

// lanEventHandler is used to handle events from the lan Serf cluster
func (s *Server) lanEventHandler() {
	for {
		select {
		case e := <-s.eventChLAN:
			s.logger.Printf("[INFO] LAN Event: %v", e)
		case <-s.shutdownCh:
			return
		}
	}
}

// wanEventHandler is used to handle events from the wan Serf cluster
func (s *Server) wanEventHandler() {
	for {
		select {
		case e := <-s.eventChWAN:
			s.logger.Printf("[INFO] WAN Event: %v", e)
		case <-s.shutdownCh:
			return
		}
	}
}
<commit_msg>consul: type safe event handlers<commit_after>package consul

import (
	"github.com/hashicorp/serf/serf"
)

// lanEventHandler is used to handle events from the lan Serf cluster
func (s *Server) lanEventHandler() {
	for {
		select {
		case e := <-s.eventChLAN:
			switch e.EventType() {
			case serf.EventMemberJoin:
				s.localJoin(e.(serf.MemberEvent))
			case serf.EventMemberLeave:
				s.localLeave(e.(serf.MemberEvent))
			case serf.EventMemberFailed:
				s.localFailed(e.(serf.MemberEvent))
			case serf.EventUser:
				s.localEvent(e.(serf.UserEvent))
			default:
				s.logger.Printf("[WARN] Unhandled LAN Serf Event: %#v", e)
			}

		case <-s.shutdownCh:
			return
		}
	}
}

// wanEventHandler is used to handle events from the wan Serf cluster
func (s *Server) wanEventHandler() {
	for {
		select {
		case e := <-s.eventChWAN:
			switch e.EventType() {
			case serf.EventMemberJoin:
				s.remoteJoin(e.(serf.MemberEvent))
			case serf.EventMemberLeave:
				s.remoteLeave(e.(serf.MemberEvent))
			case serf.EventMemberFailed:
				s.remoteFailed(e.(serf.MemberEvent))
			case serf.EventUser:
				s.remoteEvent(e.(serf.UserEvent))
			default:
				s.logger.Printf("[WARN] Unhandled LAN Serf Event: %#v", e)
			}

		case <-s.shutdownCh:
			return
		}
	}
}

// localJoin is used to handle join events on the lan serf cluster
func (s *Server) localJoin(me serf.MemberEvent) {
}

// localLeave is used to handle leave events on the lan serf cluster
func (s *Server) localLeave(me serf.MemberEvent) {
}

// localFailed is used to handle fail events on the lan serf cluster
func (s *Server) localFailed(me serf.MemberEvent) {
}

// localEvent is used to handle events on the lan serf cluster
func (s *Server) localEvent(ue serf.UserEvent) {
}

// remoteJoin is used to handle join events on the wan serf cluster
func (s *Server) remoteJoin(me serf.MemberEvent) {
}

// remoteLeave is used to handle leave events on the wan serf cluster
func (s *Server) remoteLeave(me serf.MemberEvent) {
}

// remoteFailed is used to handle fail events on the wan serf cluster
func (s *Server) remoteFailed(me serf.MemberEvent) {
}

// remoteEvent is used to handle events on the wan serf cluster
func (s *Server) remoteEvent(ue serf.UserEvent) {
}
<|endoftext|>


<commit_before>package kontrol

import (
	"errors"
	"fmt"
	"strings"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/koding/kite"
	"github.com/koding/kite/dnode"
	"github.com/koding/kite/kitekey"
	"github.com/koding/kite/kontrol/onceevery"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
)

func (k *Kontrol) handleRegister(r *kite.Request) (interface{}, error) {
	k.log.Info("Register request from: %s", r.Client.Kite)

	if r.Args.One().MustMap()["url"].MustString() == "" {
		return nil, errors.New("invalid url")
	}

	var args struct {
		URL string `json:"url"`
	}
	r.Args.One().MustUnmarshal(&args)
	if args.URL == "" {
		return nil, errors.New("empty url")
	}

	// Only accept requests with kiteKey because we need this info
	// for generating tokens for this kite.
	if r.Auth.Type != "kiteKey" {
		return nil, fmt.Errorf("Unexpected authentication type: %s", r.Auth.Type)
	}

	t, err := jwt.Parse(r.Auth.Key, kitekey.GetKontrolKey)
	if err != nil {
		return nil, err
	}

	publicKey, ok := t.Claims["kontrolKey"].(string)
	if !ok {
		return nil, errors.New("public key is not passed")
	}

	var keyPair *KeyPair
	var newKey bool

	// check if the key is valid and is stored in the key pair storage, if not
	// check if there is a new key we can use.
	keyPair, err = k.keyPair.GetKeyFromPublic(strings.TrimSpace(publicKey))
	if err != nil {
		newKey = true
		keyPair, err = k.pickKey(r)
		if err != nil {
			return nil, err // nothing to do here ..
		}
	}

	kiteURL := args.URL
	remote := r.Client

	if err := validateKiteKey(&remote.Kite); err != nil {
		return nil, err
	}

	value := &kontrolprotocol.RegisterValue{
		URL:   kiteURL,
		KeyID: keyPair.ID,
	}

	// Register first by adding the value to the storage. Return if there is
	// any error.
	if err := k.storage.Upsert(&remote.Kite, value); err != nil {
		k.log.Error("storage add '%s' error: %s", remote.Kite, err)
		return nil, errors.New("internal error - register")
	}

	every := onceevery.New(UpdateInterval)

	ping := make(chan struct{}, 1)
	closed := false

	updaterFunc := func() {
		for {
			select {
			case <-ping:
				k.log.Debug("Kite is active, got a ping %s", remote.Kite)
				every.Do(func() {
					k.log.Debug("Kite is active, updating the value %s", remote.Kite)
					err := k.storage.Update(&remote.Kite, value)
					if err != nil {
						k.log.Error("storage update '%s' error: %s", remote.Kite, err)
					}
				})
			case <-time.After(HeartbeatInterval + HeartbeatDelay):
				k.log.Debug("Kite didn't sent any heartbeat %s.", remote.Kite)
				every.Stop()
				closed = true
				return
			}
		}
	}
	go updaterFunc()

	heartbeatArgs := []interface{}{
		HeartbeatInterval / time.Second,
		dnode.Callback(func(args *dnode.Partial) {
			k.log.Debug("Kite send us an heartbeat. %s", remote.Kite)

			k.clientLocks.Get(remote.Kite.ID).Lock()
			defer k.clientLocks.Get(remote.Kite.ID).Unlock()

			select {
			case ping <- struct{}{}:
			default:
			}

			// seems we miss a heartbeat, so start it again!
			if closed {
				closed = false
				k.log.Warning("Updater was closed, but we are still getting heartbeats. Starting again %s",
					remote.Kite)

				// it might be removed because the ttl cleaner would come
				// before us, so try to add it again, the updater will than
				// continue to update it afterwards.
				k.storage.Upsert(&remote.Kite, value)
				go updaterFunc()
			}
		}),
	}

	// now trigger the remote kite so it sends us periodically an heartbeat
	remote.GoWithTimeout("kite.heartbeat", 4*time.Second, heartbeatArgs...)

	k.log.Info("Kite registered: %s", remote.Kite)

	remote.OnDisconnect(func() {
		k.log.Info("Kite disconnected: %s", remote.Kite)
		every.Stop()
	})

	// send response back to the kite, also send the new public Key if it's exist
	p := &protocol.RegisterResult{URL: args.URL}
	if newKey {
		p.PublicKey = keyPair.Public
	}

	return p, nil
}

func (k *Kontrol) handleGetKites(r *kite.Request) (interface{}, error) {
	var args protocol.GetKitesArgs
	if err := r.Args.One().Unmarshal(&args); err != nil {
		return nil, err
	}

	// Get kites from the storage
	kites, err := k.storage.Get(args.Query)
	if err != nil {
		return nil, err
	}

	for _, kite := range kites {
		// audience will go into the token as "aud" claim.
		audience := getAudience(args.Query)

		keyPair, err := k.keyPair.GetKeyFromID(kite.KeyID)
		if err != nil {
			return nil, err
		}

		// Generate token once here because we are using the same token for every
		// kite we return and generating many tokens is really slow.
		token, err := generateToken(audience, r.Username,
			k.Kite.Kite().Username, keyPair.Private)
		if err != nil {
			return nil, err
		}

		kite.Token = token
	}

	return &protocol.GetKitesResult{
		Kites: kites,
	}, nil
}

func (k *Kontrol) handleGetToken(r *kite.Request) (interface{}, error) {
	var query *protocol.KontrolQuery
	err := r.Args.One().Unmarshal(&query)
	if err != nil {
		return nil, errors.New("Invalid query")
	}

	// check if it's exist
	kites, err := k.storage.Get(query)
	if err != nil {
		return nil, err
	}

	if len(kites) > 1 {
		return nil, errors.New("query matches more than one kite")
	}

	kite := kites[0]
	audience := getAudience(query)

	keyPair, err := k.keyPair.GetKeyFromID(kite.KeyID)
	if err != nil {
		return nil, err
	}

	return generateToken(audience, r.Username, k.Kite.Kite().Username, keyPair.Private)
}

func (k *Kontrol) handleMachine(r *kite.Request) (interface{}, error) {
	var args struct {
		AuthType string
	}

	err := r.Args.One().Unmarshal(&args)
	if err != nil {
		return nil, err
	}

	if k.MachineAuthenticate != nil {
		// an empty authType is ok, the implementer is responsible of it. It
		// can care of it or it can return an error
		if err := k.MachineAuthenticate(args.AuthType, r); err != nil {
			return nil, errors.New("cannot authenticate user")
		}
	}

	keyPair, err := k.pickKey(r)
	if err != nil {
		return nil, err
	}

	return k.registerUser(r.Client.Kite.Username, keyPair.Public, keyPair.Private)
}

func (k *Kontrol) handleGetKey(r *kite.Request) (interface{}, error) {
	// Only accept requests with kiteKey because we need this info
	// for checking if the key is valid and needs to be regenerated
	if r.Auth.Type != "kiteKey" {
		return nil, fmt.Errorf("Unexpected authentication type: %s", r.Auth.Type)
	}

	t, err := jwt.Parse(r.Auth.Key, kitekey.GetKontrolKey)
	if err != nil {
		return nil, err
	}

	publicKey, ok := t.Claims["kontrolKey"].(string)
	if !ok {
		return nil, errors.New("public key is not passed")
	}

	err = k.keyPair.IsValid(publicKey)
	if err == nil {
		// everything is ok, just return the old one
		return publicKey, nil
	}

	keyPair, err := k.pickKey(r)
	if err != nil {
		return nil, err
	}

	return keyPair.Public, nil
}

func (k *Kontrol) pickKey(r *kite.Request) (*KeyPair, error) {
	if k.MachineKeyPicker != nil {
		keyPair, err := k.MachineKeyPicker(r)
		if err != nil {
			return nil, err
		}
		return keyPair, nil
	}

	if len(k.lastPublic) != 0 && len(k.lastPrivate) != 0 {
		return &KeyPair{
			Public:  k.lastPublic[len(k.lastPublic)-1],
			Private: k.lastPrivate[len(k.lastPrivate)-1],
			ID:      k.lastIDs[len(k.lastIDs)-1],
		}, nil
	}

	k.log.Error("neither machineKeyPicker neither public/private keys are available")
	return nil, errors.New("internal error - 1")
}
<commit_msg>handlers: check for zero len kites result<commit_after>package kontrol

import (
	"errors"
	"fmt"
	"strings"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/koding/kite"
	"github.com/koding/kite/dnode"
	"github.com/koding/kite/kitekey"
	"github.com/koding/kite/kontrol/onceevery"
	kontrolprotocol "github.com/koding/kite/kontrol/protocol"
	"github.com/koding/kite/protocol"
)

func (k *Kontrol) handleRegister(r *kite.Request) (interface{}, error) {
	k.log.Info("Register request from: %s", r.Client.Kite)

	if r.Args.One().MustMap()["url"].MustString() == "" {
		return nil, errors.New("invalid url")
	}

	var args struct {
		URL string `json:"url"`
	}
	r.Args.One().MustUnmarshal(&args)
	if args.URL == "" {
		return nil, errors.New("empty url")
	}

	// Only accept requests with kiteKey because we need this info
	// for generating tokens for this kite.
	if r.Auth.Type != "kiteKey" {
		return nil, fmt.Errorf("Unexpected authentication type: %s", r.Auth.Type)
	}

	t, err := jwt.Parse(r.Auth.Key, kitekey.GetKontrolKey)
	if err != nil {
		return nil, err
	}

	publicKey, ok := t.Claims["kontrolKey"].(string)
	if !ok {
		return nil, errors.New("public key is not passed")
	}

	var keyPair *KeyPair
	var newKey bool

	// check if the key is valid and is stored in the key pair storage, if not
	// check if there is a new key we can use.
	keyPair, err = k.keyPair.GetKeyFromPublic(strings.TrimSpace(publicKey))
	if err != nil {
		newKey = true
		keyPair, err = k.pickKey(r)
		if err != nil {
			return nil, err // nothing to do here ..
		}
	}

	kiteURL := args.URL
	remote := r.Client

	if err := validateKiteKey(&remote.Kite); err != nil {
		return nil, err
	}

	value := &kontrolprotocol.RegisterValue{
		URL:   kiteURL,
		KeyID: keyPair.ID,
	}

	// Register first by adding the value to the storage. Return if there is
	// any error.
	if err := k.storage.Upsert(&remote.Kite, value); err != nil {
		k.log.Error("storage add '%s' error: %s", remote.Kite, err)
		return nil, errors.New("internal error - register")
	}

	every := onceevery.New(UpdateInterval)

	ping := make(chan struct{}, 1)
	closed := false

	updaterFunc := func() {
		for {
			select {
			case <-ping:
				k.log.Debug("Kite is active, got a ping %s", remote.Kite)
				every.Do(func() {
					k.log.Debug("Kite is active, updating the value %s", remote.Kite)
					err := k.storage.Update(&remote.Kite, value)
					if err != nil {
						k.log.Error("storage update '%s' error: %s", remote.Kite, err)
					}
				})
			case <-time.After(HeartbeatInterval + HeartbeatDelay):
				k.log.Debug("Kite didn't sent any heartbeat %s.", remote.Kite)
				every.Stop()
				closed = true
				return
			}
		}
	}
	go updaterFunc()

	heartbeatArgs := []interface{}{
		HeartbeatInterval / time.Second,
		dnode.Callback(func(args *dnode.Partial) {
			k.log.Debug("Kite send us an heartbeat. %s", remote.Kite)

			k.clientLocks.Get(remote.Kite.ID).Lock()
			defer k.clientLocks.Get(remote.Kite.ID).Unlock()

			select {
			case ping <- struct{}{}:
			default:
			}

			// seems we miss a heartbeat, so start it again!
			if closed {
				closed = false
				k.log.Warning("Updater was closed, but we are still getting heartbeats. Starting again %s",
					remote.Kite)

				// it might be removed because the ttl cleaner would come
				// before us, so try to add it again, the updater will than
				// continue to update it afterwards.
				k.storage.Upsert(&remote.Kite, value)
				go updaterFunc()
			}
		}),
	}

	// now trigger the remote kite so it sends us periodically an heartbeat
	remote.GoWithTimeout("kite.heartbeat", 4*time.Second, heartbeatArgs...)

	k.log.Info("Kite registered: %s", remote.Kite)

	remote.OnDisconnect(func() {
		k.log.Info("Kite disconnected: %s", remote.Kite)
		every.Stop()
	})

	// send response back to the kite, also send the new public Key if it's exist
	p := &protocol.RegisterResult{URL: args.URL}
	if newKey {
		p.PublicKey = keyPair.Public
	}

	return p, nil
}

func (k *Kontrol) handleGetKites(r *kite.Request) (interface{}, error) {
	var args protocol.GetKitesArgs
	if err := r.Args.One().Unmarshal(&args); err != nil {
		return nil, err
	}

	// Get kites from the storage
	kites, err := k.storage.Get(args.Query)
	if err != nil {
		return nil, err
	}

	for _, kite := range kites {
		// audience will go into the token as "aud" claim.
		audience := getAudience(args.Query)

		keyPair, err := k.keyPair.GetKeyFromID(kite.KeyID)
		if err != nil {
			return nil, err
		}

		// Generate token once here because we are using the same token for every
		// kite we return and generating many tokens is really slow.
		token, err := generateToken(audience, r.Username,
			k.Kite.Kite().Username, keyPair.Private)
		if err != nil {
			return nil, err
		}

		kite.Token = token
	}

	return &protocol.GetKitesResult{
		Kites: kites,
	}, nil
}

func (k *Kontrol) handleGetToken(r *kite.Request) (interface{}, error) {
	var query *protocol.KontrolQuery
	err := r.Args.One().Unmarshal(&query)
	if err != nil {
		return nil, errors.New("Invalid query")
	}

	// check if it's exist
	kites, err := k.storage.Get(query)
	if err != nil {
		return nil, err
	}

	if len(kites) > 1 {
		return nil, errors.New("query matches more than one kite")
	}

	if len(kites) == 0 {
		return nil, errors.New("no kites found")
	}

	kite := kites[0]
	audience := getAudience(query)

	keyPair, err := k.keyPair.GetKeyFromID(kite.KeyID)
	if err != nil {
		return nil, err
	}

	return generateToken(audience, r.Username, k.Kite.Kite().Username, keyPair.Private)
}

func (k *Kontrol) handleMachine(r *kite.Request) (interface{}, error) {
	var args struct {
		AuthType string
	}

	err := r.Args.One().Unmarshal(&args)
	if err != nil {
		return nil, err
	}

	if k.MachineAuthenticate != nil {
		// an empty authType is ok, the implementer is responsible of it. It
		// can care of it or it can return an error
		if err := k.MachineAuthenticate(args.AuthType, r); err != nil {
			return nil, errors.New("cannot authenticate user")
		}
	}

	keyPair, err := k.pickKey(r)
	if err != nil {
		return nil, err
	}

	return k.registerUser(r.Client.Kite.Username, keyPair.Public, keyPair.Private)
}

func (k *Kontrol) handleGetKey(r *kite.Request) (interface{}, error) {
	// Only accept requests with kiteKey because we need this info
	// for checking if the key is valid and needs to be regenerated
	if r.Auth.Type != "kiteKey" {
		return nil, fmt.Errorf("Unexpected authentication type: %s", r.Auth.Type)
	}

	t, err := jwt.Parse(r.Auth.Key, kitekey.GetKontrolKey)
	if err != nil {
		return nil, err
	}

	publicKey, ok := t.Claims["kontrolKey"].(string)
	if !ok {
		return nil, errors.New("public key is not passed")
	}

	err = k.keyPair.IsValid(publicKey)
	if err == nil {
		// everything is ok, just return the old one
		return publicKey, nil
	}

	keyPair, err := k.pickKey(r)
	if err != nil {
		return nil, err
	}

	return keyPair.Public, nil
}

func (k *Kontrol) pickKey(r *kite.Request) (*KeyPair, error) {
	if k.MachineKeyPicker != nil {
		keyPair, err := k.MachineKeyPicker(r)
		if err != nil {
			return nil, err
		}
		return keyPair, nil
	}

	if len(k.lastPublic) != 0 && len(k.lastPrivate) != 0 {
		return &KeyPair{
			Public:  k.lastPublic[len(k.lastPublic)-1],
			Private: k.lastPrivate[len(k.lastPrivate)-1],
			ID:      k.lastIDs[len(k.lastIDs)-1],
		}, nil
	}

	k.log.Error("neither machineKeyPicker neither public/private keys are available")
	return nil, errors.New("internal error - 1")
}
<|endoftext|>


<commit_before>package morningStar

import (
	"../jsonHttp"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"sync"
	"time"
)

const PERFORMANCE_URL = `http://www.morningstar.fr/fr/funds/snapshot/snapshot.aspx?tab=1&id=`
const VOLATILITE_URL = `http://www.morningstar.fr/fr/funds/snapshot/snapshot.aspx?tab=2&id=`
const REFRESH_DELAY = 18
const CONCURRENT_FETCHER = 20

var EMPTY_BYTE = []byte(``)
var ZERO_BYTE = []byte(`0`)
var PERIOD_BYTE = []byte(`.`)
var COMMA_BYTE = []byte(`,`)
var PERCENT_BYTE = []byte(`%`)
var AMP_BYTE = []byte(`&`)
var HTML_AMP_BYTE = []byte(`&`)

var LIST_REQUEST = regexp.MustCompile(`^/list$`)
var PERF_REQUEST = regexp.MustCompile(`^/(.+?)$`)

var ISIN = regexp.MustCompile(`ISIN.:(\S+)`)
var LABEL = regexp.MustCompile(`<h1[^>]*?>((?:.|\n)*?)</h1>`)
var RATING = regexp.MustCompile(`<span\sclass=".*?stars([0-9]).*?">`)
var CATEGORY = regexp.MustCompile(`<span[^>]*?>Catégorie</span>.*?<span[^>]*?>(.*?)</span>`)
var PERF_ONE_MONTH = regexp.MustCompile(`<td[^>]*?>1 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_THREE_MONTH = regexp.MustCompile(`<td[^>]*?>3 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_SIX_MONTH = regexp.MustCompile(`<td[^>]*?>6 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_ONE_YEAR = regexp.MustCompile(`<td[^>]*?>1 an</td><td[^>]*?>(.*?)</td>`)
var VOL_3_YEAR = regexp.MustCompile(`<td[^>]*?>Ecart-type 3 ans.?</td><td[^>]*?>(.*?)</td>`)

type SyncedMap struct {
	sync.RWMutex
	performances map[string]Performance
}

func (m *SyncedMap) get(key string) (Performance, bool) {
	m.RLock()
	defer m.RUnlock()

	performance, ok := m.performances[key]
	return performance, ok
}

func (m *SyncedMap) push(key string, performance Performance) {
	m.Lock()
	defer m.Unlock()

	m.performances[key] = performance
}

var PERFORMANCE_CACHE = SyncedMap{performances: make(map[string]Performance)}

type Performance struct {
	Id            string    `json:"id"`
	Isin          string    `json:"isin"`
	Label         string    `json:"label"`
	Category      string    `json:"category"`
	Rating        string    `json:"rating"`
	OneMonth      float64   `json:"1m"`
	ThreeMonth    float64   `json:"3m"`
	SixMonth      float64   `json:"6m"`
	OneYear       float64   `json:"1y"`
	VolThreeYears float64   `json:"v3y"`
	Score         float64   `json:"score"`
	Update        time.Time `json:"ts"`
}

type Results struct {
	Results interface{} `json:"results"`
}

func readBody(body io.ReadCloser) ([]byte, error) {
	defer body.Close()
	return ioutil.ReadAll(body)
}

func getBody(url string) ([]byte, error) {
	response, err := http.Get(url)
	if err != nil {
		return nil, errors.New(`Error while retrieving data from ` + url)
	}

	if response.StatusCode >= 400 {
		return nil, errors.New(`Got error ` + strconv.Itoa(response.StatusCode) + ` while getting ` + url)
	}

	body, err := readBody(response.Body)
	if err != nil {
		return nil, errors.New(`Error while reading body of ` + url)
	}

	return body, nil
}

func getLabel(extract *regexp.Regexp, body []byte, defaultValue []byte) []byte {
	match := extract.FindSubmatch(body)
	if match == nil {
		return defaultValue
	}

	return bytes.Replace(match[1], HTML_AMP_BYTE, AMP_BYTE, -1)
}

func getPerformance(extract *regexp.Regexp, body []byte) float64 {
	dotResult := bytes.Replace(getLabel(extract, body, EMPTY_BYTE), COMMA_BYTE, PERIOD_BYTE, -1)
	percentageResult := bytes.Replace(dotResult, PERCENT_BYTE, EMPTY_BYTE, -1)
	trimResult := bytes.TrimSpace(percentageResult)

	result, err := strconv.ParseFloat(string(trimResult), 64)
	if err != nil {
		return 0.0
	}
	return result
}

func SinglePerformance(morningStarId []byte) (*Performance, error) {
	cleanId := string(bytes.ToLower(morningStarId))

	performance, ok := PERFORMANCE_CACHE.get(cleanId)

	if ok && time.Now().Add(time.Hour*-REFRESH_DELAY).Before(performance.Update) {
		return &performance, nil
	}

	performanceBody, err := getBody(PERFORMANCE_URL + cleanId)
	if err != nil {
		return nil, err
	}

	volatiliteBody, err := getBody(VOLATILITE_URL + cleanId)
	if err != nil {
		return nil, err
	}

	isin := string(getLabel(ISIN, performanceBody, EMPTY_BYTE))
	label := string(getLabel(LABEL, performanceBody, EMPTY_BYTE))
	rating := string(getLabel(RATING, performanceBody, ZERO_BYTE))
	category := string(getLabel(CATEGORY, performanceBody, EMPTY_BYTE))
	oneMonth := getPerformance(PERF_ONE_MONTH, performanceBody)
	threeMonths := getPerformance(PERF_THREE_MONTH, performanceBody)
	sixMonths := getPerformance(PERF_SIX_MONTH, performanceBody)
	oneYear := getPerformance(PERF_ONE_YEAR, performanceBody)
	volThreeYears := getPerformance(VOL_3_YEAR, volatiliteBody)

	score := (0.25 * oneMonth) + (0.3 * threeMonths) + (0.25 * sixMonths) + (0.2 * oneYear) - (0.1 * volThreeYears)
	scoreTruncated := float64(int(score*100)) / 100

	performance = Performance{cleanId, isin, label, category, rating, oneMonth, threeMonths, sixMonths, oneYear, volThreeYears, scoreTruncated, time.Now()}

	PERFORMANCE_CACHE.push(cleanId, performance)

	return &performance, nil
}

func singlePerformanceHandler(w http.ResponseWriter, morningStarId []byte) {
	performance, err := SinglePerformance(morningStarId)

	if err != nil {
		http.Error(w, err.Error(), 500)
	} else {
		jsonHttp.ResponseJson(w, *performance)
	}
}

func listHandler(w http.ResponseWriter, r *http.Request) {
	listBody, err := readBody(r.Body)
	if err != nil {
		http.Error(w, `Error while reading body for list`, 500)
		return
	}

	if len(bytes.TrimSpace(listBody)) == 0 {
		jsonHttp.ResponseJson(w, Results{[0]Performance{}})
		return
	}

	ids := bytes.Split(listBody, COMMA_BYTE)
	size := len(ids)

	var wg sync.WaitGroup
	wg.Add(size)

	performances := make(chan *Performance, CONCURRENT_FETCHER)
	tokens := make(chan struct{}, CONCURRENT_FETCHER)

	go func() {
		for _, id := range ids {
			tokens <- struct{}{}

			go func(morningStarId []byte) {
				defer wg.Done()
				if performance, err := SinglePerformance(morningStarId); err == nil {
					performances <- performance
				}
			}(id)

			<-tokens
		}
	}()

	go func() {
		wg.Wait()
		close(performances)
	}()

	results := make([]*Performance, 0, size)
	for performance := range performances {
		results = append(results, performance)
	}

	jsonHttp.ResponseJson(w, Results{results})
}

type Handler struct {
}

func (handler Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Add(`Access-Control-Allow-Origin`, `*`)
	w.Header().Add(`Access-Control-Allow-Headers`, `Content-Type`)
	w.Header().Add(`Access-Control-Allow-Methods`, `GET, POST`)
	w.Header().Add(`X-Content-Type-Options`, `nosniff`)

	urlPath := []byte(r.URL.Path)

	if LIST_REQUEST.Match(urlPath) {
		listHandler(w, r)
	} else if PERF_REQUEST.Match(urlPath) {
		singlePerformanceHandler(w, PERF_REQUEST.FindSubmatch(urlPath)[1])
	}
}
<commit_msg>Adding function for clarity<commit_after>package morningStar

import (
	"../jsonHttp"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"sync"
	"time"
)

const PERFORMANCE_URL = `http://www.morningstar.fr/fr/funds/snapshot/snapshot.aspx?tab=1&id=`
const VOLATILITE_URL = `http://www.morningstar.fr/fr/funds/snapshot/snapshot.aspx?tab=2&id=`
const REFRESH_DELAY = 18
const CONCURRENT_FETCHER = 20

var EMPTY_BYTE = []byte(``)
var ZERO_BYTE = []byte(`0`)
var PERIOD_BYTE = []byte(`.`)
var COMMA_BYTE = []byte(`,`)
var PERCENT_BYTE = []byte(`%`)
var AMP_BYTE = []byte(`&`)
var HTML_AMP_BYTE = []byte(`&`)

var LIST_REQUEST = regexp.MustCompile(`^/list$`)
var PERF_REQUEST = regexp.MustCompile(`^/(.+?)$`)

var ISIN = regexp.MustCompile(`ISIN.:(\S+)`)
var LABEL = regexp.MustCompile(`<h1[^>]*?>((?:.|\n)*?)</h1>`)
var RATING = regexp.MustCompile(`<span\sclass=".*?stars([0-9]).*?">`)
var CATEGORY = regexp.MustCompile(`<span[^>]*?>Catégorie</span>.*?<span[^>]*?>(.*?)</span>`)
var PERF_ONE_MONTH = regexp.MustCompile(`<td[^>]*?>1 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_THREE_MONTH = regexp.MustCompile(`<td[^>]*?>3 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_SIX_MONTH = regexp.MustCompile(`<td[^>]*?>6 mois</td><td[^>]*?>(.*?)</td>`)
var PERF_ONE_YEAR = regexp.MustCompile(`<td[^>]*?>1 an</td><td[^>]*?>(.*?)</td>`)
var VOL_3_YEAR = regexp.MustCompile(`<td[^>]*?>Ecart-type 3 ans.?</td><td[^>]*?>(.*?)</td>`)

type SyncedMap struct {
	sync.RWMutex
	performances map[string]Performance
}

func (m *SyncedMap) get(key string) (Performance, bool) {
	m.RLock()
	defer m.RUnlock()

	performance, ok := m.performances[key]
	return performance, ok
}

func (m *SyncedMap) push(key string, performance Performance) {
	m.Lock()
	defer m.Unlock()

	m.performances[key] = performance
}

var PERFORMANCE_CACHE = SyncedMap{performances: make(map[string]Performance)}

type Performance struct {
	Id            string    `json:"id"`
	Isin          string    `json:"isin"`
	Label         string    `json:"label"`
	Category      string    `json:"category"`
	Rating        string    `json:"rating"`
	OneMonth      float64   `json:"1m"`
	ThreeMonth    float64   `json:"3m"`
	SixMonth      float64   `json:"6m"`
	OneYear       float64   `json:"1y"`
	VolThreeYears float64   `json:"v3y"`
	Score         float64   `json:"score"`
	Update        time.Time `json:"ts"`
}

type Results struct {
	Results interface{} `json:"results"`
}

func readBody(body io.ReadCloser) ([]byte, error) {
	defer body.Close()
	return ioutil.ReadAll(body)
}

func getBody(url string) ([]byte, error) {
	response, err := http.Get(url)
	if err != nil {
		return nil, errors.New(`Error while retrieving data from ` + url)
	}

	if response.StatusCode >= 400 {
		return nil, errors.New(`Got error ` + strconv.Itoa(response.StatusCode) + ` while getting ` + url)
	}

	body, err := readBody(response.Body)
	if err != nil {
		return nil, errors.New(`Error while reading body of ` + url)
	}

	return body, nil
}

func getLabel(extract *regexp.Regexp, body []byte, defaultValue []byte) []byte {
	match := extract.FindSubmatch(body)
	if match == nil {
		return defaultValue
	}

	return bytes.Replace(match[1], HTML_AMP_BYTE, AMP_BYTE, -1)
}

func getPerformance(extract *regexp.Regexp, body []byte) float64 {
	dotResult := bytes.Replace(getLabel(extract, body, EMPTY_BYTE), COMMA_BYTE, PERIOD_BYTE, -1)
	percentageResult := bytes.Replace(dotResult, PERCENT_BYTE, EMPTY_BYTE, -1)
	trimResult := bytes.TrimSpace(percentageResult)

	result, err := strconv.ParseFloat(string(trimResult), 64)
	if err != nil {
		return 0.0
	}
	return result
}

func SinglePerformance(morningStarId []byte) (*Performance, error) {
	cleanId := string(bytes.ToLower(morningStarId))

	performance, ok := PERFORMANCE_CACHE.get(cleanId)

	if ok && time.Now().Add(time.Hour*-REFRESH_DELAY).Before(performance.Update) {
		return &performance, nil
	}

	performanceBody, err := getBody(PERFORMANCE_URL + cleanId)
	if err != nil {
		return nil, err
	}

	volatiliteBody, err := getBody(VOLATILITE_URL + cleanId)
	if err != nil {
		return nil, err
	}

	isin := string(getLabel(ISIN, performanceBody, EMPTY_BYTE))
	label := string(getLabel(LABEL, performanceBody, EMPTY_BYTE))
	rating := string(getLabel(RATING, performanceBody, ZERO_BYTE))
	category := string(getLabel(CATEGORY, performanceBody, EMPTY_BYTE))
	oneMonth := getPerformance(PERF_ONE_MONTH, performanceBody)
	threeMonths := getPerformance(PERF_THREE_MONTH, performanceBody)
	sixMonths := getPerformance(PERF_SIX_MONTH, performanceBody)
	oneYear := getPerformance(PERF_ONE_YEAR, performanceBody)
	volThreeYears := getPerformance(VOL_3_YEAR, volatiliteBody)

	score := (0.25 * oneMonth) + (0.3 * threeMonths) + (0.25 * sixMonths) + (0.2 * oneYear) - (0.1 * volThreeYears)
	scoreTruncated := float64(int(score*100)) / 100

	performance = Performance{cleanId, isin, label, category, rating, oneMonth, threeMonths, sixMonths, oneYear, volThreeYears, scoreTruncated, time.Now()}

	PERFORMANCE_CACHE.push(cleanId, performance)

	return &performance, nil
}

func singlePerformanceHandler(w http.ResponseWriter, morningStarId []byte) {
	performance, err := SinglePerformance(morningStarId)

	if err != nil {
		http.Error(w, err.Error(), 500)
	} else {
		jsonHttp.ResponseJson(w, *performance)
	}
}

func allPerformances(ids [][]byte, wg sync.WaitGroup, performances chan<- *Performance) {
	tokens := make(chan struct{}, CONCURRENT_FETCHER)
	wg.Add(len(ids))

	for _, id := range ids {
		tokens <- struct{}{}

		go func(morningStarId []byte) {
			defer wg.Done()
			if performance, err := SinglePerformance(morningStarId); err == nil {
				performances <- performance
			}
		}(id)

		<-tokens
	}
}

func listHandler(w http.ResponseWriter, r *http.Request) {
	listBody, err := readBody(r.Body)
	if err != nil {
		http.Error(w, `Error while reading body for list`, 500)
		return
	}

	if len(bytes.TrimSpace(listBody)) == 0 {
		jsonHttp.ResponseJson(w, Results{[0]Performance{}})
		return
	}

	var wg sync.WaitGroup
	performances := make(chan *Performance, CONCURRENT_FETCHER)

	go allPerformances(bytes.Split(listBody, COMMA_BYTE), wg, performances)

	go func() {
		wg.Wait()
		close(performances)
	}()

	results := make([]*Performance, 0, len(ids))
	for performance := range performances {
		results = append(results, performance)
	}

	jsonHttp.ResponseJson(w, Results{results})
}

type Handler struct {
}

func (handler Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Add(`Access-Control-Allow-Origin`, `*`)
	w.Header().Add(`Access-Control-Allow-Headers`, `Content-Type`)
	w.Header().Add(`Access-Control-Allow-Methods`, `GET, POST`)
	w.Header().Add(`X-Content-Type-Options`, `nosniff`)

	urlPath := []byte(r.URL.Path)

	if LIST_REQUEST.Match(urlPath) {
		listHandler(w, r)
	} else if PERF_REQUEST.Match(urlPath) {
		singlePerformanceHandler(w, PERF_REQUEST.FindSubmatch(urlPath)[1])
	}
}
<|endoftext|>


<commit_before>// @author Couchbase <info@couchbase.com>
// @copyright 2014 Couchbase, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cbauth

import (
	"encoding/json"
	"net/http"
)

// SetRequestAuthVia sets basic auth header in given http request
// according to given authenticator. It will extract target
// hostname/port from request and figure out right service credentials
// for that endpoint. If nil authenticator is passed, Default
// authenticator is used.
func SetRequestAuthVia(req *http.Request, a Authenticator) error {
	return WithAuthenticator(a, func(a Authenticator) (err error) {
		user, pwd, err := a.GetHTTPServiceAuth(req.URL.Host)
		if err != nil {
			return
		}
		req.SetBasicAuth(user, pwd)
		return
	})
}

// SetRequestAuth sets basic auth header in given http request
// according to default authenticator. Simply calls SetRequestAuthVia
// with nil authenticator.
func SetRequestAuth(req *http.Request) error {
	return SetRequestAuthVia(req, nil)
}

func duplicateStringsSlice(in []string) []string {
	return append([]string{}, in...)
}

func dupHeaders(h http.Header) http.Header {
	rv := make(http.Header)
	for k, v := range h {
		rv[k] = duplicateStringsSlice(v)
	}
	return rv
}

func dupRequest(req *http.Request) *http.Request {
	rv := *req
	rv.Header = dupHeaders(req.Header)
	rv.Trailer = dupHeaders(req.Trailer)
	return &rv
}

type cbauthRoundTripper struct {
	slave http.RoundTripper
	a     Authenticator
}

func (rt *cbauthRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	req = dupRequest(req)
	if err := SetRequestAuthVia(req, rt.a); err != nil {
		return nil, err
	}
	return rt.slave.RoundTrip(req)
}

// WrapHTTPTransport constructs http transport that automatically does
// SetRequestAuthVia for requests it sends. As usual, if nil
// authenticator is passed, default authenticator is used.
func WrapHTTPTransport(transport http.RoundTripper, a Authenticator) http.RoundTripper {
	return &cbauthRoundTripper{
		slave: transport,
		a:     a,
	}
}

// SendUnauthorized sends 401 Unauthorized response on given response writer.
func SendUnauthorized(w http.ResponseWriter) {
	w.Header().Set("WWW-Authenticate", "Basic realm=\"Couchbase\"")
	http.Error(w, "need auth", http.StatusUnauthorized)
}

// ForbiddenJson returns json 403 response for given permission
func ForbiddenJson(permission string) ([]byte, error) {
	jsonStruct := map[string]interface{}{
		"message":     "Forbidden. User needs one of the following permissions",
		"permissions": [...]string{permission},
	}
	return json.Marshal(jsonStruct)
}

// SendForbidden sends 403 Forbidden with json payload that contains list
// of required permissions to response on given response writer.
func SendForbidden(w http.ResponseWriter, permission string) error {
	b, err := ForbiddenJson(permission)
	if err != nil {
		return err
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusForbidden)
	w.Write(b)
	return nil
}
<commit_msg>fix golint error<commit_after>// @author Couchbase <info@couchbase.com>
// @copyright 2014 Couchbase, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cbauth

import (
	"encoding/json"
	"net/http"
)

// SetRequestAuthVia sets basic auth header in given http request
// according to given authenticator. It will extract target
// hostname/port from request and figure out right service credentials
// for that endpoint. If nil authenticator is passed, Default
// authenticator is used.
func SetRequestAuthVia(req *http.Request, a Authenticator) error {
	return WithAuthenticator(a, func(a Authenticator) (err error) {
		user, pwd, err := a.GetHTTPServiceAuth(req.URL.Host)
		if err != nil {
			return
		}
		req.SetBasicAuth(user, pwd)
		return
	})
}

// SetRequestAuth sets basic auth header in given http request
// according to default authenticator. Simply calls SetRequestAuthVia
// with nil authenticator.
func SetRequestAuth(req *http.Request) error {
	return SetRequestAuthVia(req, nil)
}

func duplicateStringsSlice(in []string) []string {
	return append([]string{}, in...)
}

func dupHeaders(h http.Header) http.Header {
	rv := make(http.Header)
	for k, v := range h {
		rv[k] = duplicateStringsSlice(v)
	}
	return rv
}

func dupRequest(req *http.Request) *http.Request {
	rv := *req
	rv.Header = dupHeaders(req.Header)
	rv.Trailer = dupHeaders(req.Trailer)
	return &rv
}

type cbauthRoundTripper struct {
	slave http.RoundTripper
	a     Authenticator
}

func (rt *cbauthRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	req = dupRequest(req)
	if err := SetRequestAuthVia(req, rt.a); err != nil {
		return nil, err
	}
	return rt.slave.RoundTrip(req)
}

// WrapHTTPTransport constructs http transport that automatically does
// SetRequestAuthVia for requests it sends. As usual, if nil
// authenticator is passed, default authenticator is used.
func WrapHTTPTransport(transport http.RoundTripper, a Authenticator) http.RoundTripper {
	return &cbauthRoundTripper{
		slave: transport,
		a:     a,
	}
}

// SendUnauthorized sends 401 Unauthorized response on given response writer.
func SendUnauthorized(w http.ResponseWriter) {
	w.Header().Set("WWW-Authenticate", "Basic realm=\"Couchbase\"")
	http.Error(w, "need auth", http.StatusUnauthorized)
}

// ForbiddenJSON returns json 403 response for given permission
func ForbiddenJSON(permission string) ([]byte, error) {
	jsonStruct := map[string]interface{}{
		"message":     "Forbidden. User needs one of the following permissions",
		"permissions": [...]string{permission},
	}
	return json.Marshal(jsonStruct)
}

// SendForbidden sends 403 Forbidden with json payload that contains list
// of required permissions to response on given response writer.
func SendForbidden(w http.ResponseWriter, permission string) error {
	b, err := ForbiddenJSON(permission)
	if err != nil {
		return err
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusForbidden)
	w.Write(b)
	return nil
}
<|endoftext|>


<commit_before>package l10n_test

import (
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/qor/l10n"
)

type Product struct {
	ID              int    `gorm:"primary_key"`
	Code            string `l10n:"sync"`
	Name            string
	DeletedAt       *time.Time
	ColorVariations []ColorVariation
	Brand           Brand
	Tags            []Tag `gorm:"many2many:product_tags"`
	l10n.Locale
}

// func (Product) LocaleCreateable() {}

type ColorVariation struct {
	ID       int `gorm:"primary_key"`
	Quantity int
	Color    Color
}

type Color struct {
	ID   int `gorm:"primary_key"`
	Code string
	Name string
	l10n.Locale
}

type Brand struct {
	ID int `gorm:"primary_key"`
	l10n.Locale
}

type Tag struct {
	ID   int `gorm:"primary_key"`
	Name string
	l10n.Locale
}

var dbGlobal, dbCN, dbEN *gorm.DB

func init() {
	// CREATE USER 'qor'@'localhost' IDENTIFIED BY 'qor';
	// CREATE DATABASE qor_l10n;
	// GRANT ALL ON qor_l10n.* TO 'gorm'@'localhost';
	db, _ := gorm.Open("mysql", "qor:qor@/qor_l10n?charset=utf8&parseTime=True")
	l10n.RegisterCallbacks(&db)

	db.DropTable(&Product{})
	db.DropTable(&Tag{})
	db.Exec("drop table product_tags;")
	db.AutoMigrate(&Product{})
	db.AutoMigrate(&Tag{})

	dbGlobal = &db
	dbCN = dbGlobal.Set("l10n:locale", "zh")
	dbEN = dbGlobal.Set("l10n:locale", "en")
}
<commit_msg>Make l10n test db support env variable for db user and password.<commit_after>package l10n_test

import (
	"fmt"
	"os"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/qor/l10n"
)

type Product struct {
	ID              int    `gorm:"primary_key"`
	Code            string `l10n:"sync"`
	Name            string
	DeletedAt       *time.Time
	ColorVariations []ColorVariation
	Brand           Brand
	Tags            []Tag `gorm:"many2many:product_tags"`
	l10n.Locale
}

// func (Product) LocaleCreateable() {}

type ColorVariation struct {
	ID       int `gorm:"primary_key"`
	Quantity int
	Color    Color
}

type Color struct {
	ID   int `gorm:"primary_key"`
	Code string
	Name string
	l10n.Locale
}

type Brand struct {
	ID int `gorm:"primary_key"`
	l10n.Locale
}

type Tag struct {
	ID   int `gorm:"primary_key"`
	Name string
	l10n.Locale
}

var dbGlobal, dbCN, dbEN *gorm.DB

func init() {
	// CREATE USER 'qor'@'localhost' IDENTIFIED BY 'qor';
	// CREATE DATABASE qor_l10n;
	// GRANT ALL ON qor_l10n.* TO 'gorm'@'localhost';
	dbuser, dbpwd := "qor", "qor"
	if os.Getenv("TEST_ENV") == "CI" {
		dbuser, dbpwd = os.Getenv("DB_USER"), os.Getenv("DB_PWD")
	}
	db, _ := gorm.Open("mysql", fmt.Sprintf("%s:%s@/qor_test?charset=utf8&parseTime=True&loc=Local", dbuser, dbpwd))
	l10n.RegisterCallbacks(&db)

	db.DropTableIfExists(&Product{})
	db.DropTableIfExists(&Tag{})
	db.Exec("drop table product_tags;")
	db.AutoMigrate(&Product{})
	db.AutoMigrate(&Tag{})

	dbGlobal = &db
	dbCN = dbGlobal.Set("l10n:locale", "zh")
	dbEN = dbGlobal.Set("l10n:locale", "en")
}
<|endoftext|>


<commit_before>package cli

import (
	"bytes"
	"io/ioutil"
	"reflect"
	"strings"
	"testing"

	"github.com/elpinal/coco3/config"
	"github.com/elpinal/coco3/editor"
)

func TestFlagC(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
	}
	args := []string{"-c", "echo aaa"}
	code := c.Run(args)
	if code != 0 {
		t.Errorf("Run: got %v, want %v", code, 0)
	}
	if got, want := out.String(), "aaa\n"; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestArgs(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
	}
	args := []string{"testdata/basic.coco"}
	code := c.Run(args)
	if code != 0 {
		t.Errorf("Run: got %v, want %v", code, 0)
	}
	if got, want := out.String(), "aaa\nbbb\n"; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestExit(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
	}
	args := []string{"-c", "echo aaa; exit 42"}
	code := c.Run(args)
	if code != 42 {
		t.Errorf("Run: got %v, want %v", code, 42)
	}
	if got, want := out.String(), "aaa\n"; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestStartUpCommand(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
		Config: &config.Config{
			StartUpCommand: []byte("echo startup..."),
		},
	}
	args := []string{"-c", "echo aaa; echo bbb"}
	code := c.Run(args)
	if code != 0 {
		t.Errorf("Run: got %v, want %v", code, 0)
	}
	if got, want := out.String(), "startup...\naaa\nbbb\n"; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestExitInStartUp(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
		Config: &config.Config{
			StartUpCommand: []byte("exit 21"),
		},
	}
	args := []string{"-c", "echo aaa; exit 42"}
	code := c.Run(args)
	if code != 21 {
		t.Errorf("Run: got %v, want %v", code, 42)
	}
	if got, want := out.String(), ""; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestExitInFiles(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
	}
	args := []string{"testdata/exit1.coco", "testdata/exit2.coco"}
	code := c.Run(args)
	if code != 42 {
		t.Errorf("Run: got %v, want %v", code, 42)
	}
	golden, e := ioutil.ReadFile("testdata/exit1.golden")
	if e != nil {
		t.Errorf("reading a golden file: %v", e)
	}
	if got, want := out.String(), string(golden); got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func BenchmarkCLI(b *testing.B) {
	c := CLI{Out: ioutil.Discard}
	for i := 0; i < b.N; i++ {
		_ = c.Run([]string{"-c", "echo 1024; echo 2048"})
	}
}

func TestCompareRunes(t *testing.T) {
	tests := []struct {
		r    []rune
		s    []rune
		want bool
	}{
		{
			r:    []rune(""),
			s:    []rune(""),
			want: true,
		},
		{
			r:    []rune("a"),
			s:    []rune("a"),
			want: true,
		},
		{
			r:    []rune("abc"),
			s:    []rune("abc"),
			want: true,
		},
		{
			r:    []rune(""),
			s:    []rune("a"),
			want: false,
		},
		{
			r:    []rune("a"),
			s:    []rune("ab"),
			want: false,
		},
		{
			r:    []rune("abc"),
			s:    []rune("aba"),
			want: false,
		},
		{
			r:    []rune("aaaaa"),
			s:    []rune("aaaab"),
			want: false,
		},
	}
	for i, test := range tests {
		got := compareRunes(test.r, test.s)
		if got != test.want {
			t.Errorf("compareRunes/%d: want %v, got %v", i, test.want, got)
		}
	}
}

func TestSanitizeHistory(t *testing.T) {
	history := []string{
		"",
		"a",
		"a",
		"",
		"a",
		"b",
		"c",
		"c",
		"",
		"b",
	}
	histRunes := sanitizeHistory(history)
	want := [][]rune{
		[]rune("a"),
		[]rune("b"),
		[]rune("c"),
		[]rune("b"),
	}
	if !reflect.DeepEqual(histRunes, want) {
		t.Errorf("want %v, got %v", want, histRunes)
	}
}

func BenchmarkRun(b *testing.B) {
	for i := 0; i < b.N; i++ {
		c := CLI{
			Config: &config.Config{
				Extra: true,
			},
		}
		_ = c.Run([]string{"-c", " "})
	}
}

type benchmarkReader struct {
	ch chan struct{}
	r  *strings.Reader
}

// var errRead = errors.New("benchmarkReader: just error")

func (r *benchmarkReader) Read(p []byte) (n int, err error) {
	select {
	case r.ch <- struct{}{}:
	default:
	}
	return r.r.Read(p)
}

func TestStartup(t *testing.T) {
	c := CLI{
		In: strings.NewReader(string(editor.CharEscape) + ":q" + string(editor.CharCtrlM)),
		Config: &config.Config{
			HistFile: ":memory:",
		},
	}
	n := c.Run(nil)
	if want := 0; n != want {
		t.Errorf("Run = %d, want %d", n, want)
	}
}
<commit_msg>Benchmark start-up time<commit_after>package cli

import (
	"bytes"
	"io/ioutil"
	"reflect"
	"strings"
	"testing"

	"github.com/elpinal/coco3/config"
	"github.com/elpinal/coco3/editor"
)

func TestFlagC(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
	}
	args := []string{"-c", "echo aaa"}
	code := c.Run(args)
	if code != 0 {
		t.Errorf("Run: got %v, want %v", code, 0)
	}
	if got, want := out.String(), "aaa\n"; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestArgs(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
	}
	args := []string{"testdata/basic.coco"}
	code := c.Run(args)
	if code != 0 {
		t.Errorf("Run: got %v, want %v", code, 0)
	}
	if got, want := out.String(), "aaa\nbbb\n"; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestExit(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
	}
	args := []string{"-c", "echo aaa; exit 42"}
	code := c.Run(args)
	if code != 42 {
		t.Errorf("Run: got %v, want %v", code, 42)
	}
	if got, want := out.String(), "aaa\n"; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestStartUpCommand(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
		Config: &config.Config{
			StartUpCommand: []byte("echo startup..."),
		},
	}
	args := []string{"-c", "echo aaa; echo bbb"}
	code := c.Run(args)
	if code != 0 {
		t.Errorf("Run: got %v, want %v", code, 0)
	}
	if got, want := out.String(), "startup...\naaa\nbbb\n"; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestExitInStartUp(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
		Config: &config.Config{
			StartUpCommand: []byte("exit 21"),
		},
	}
	args := []string{"-c", "echo aaa; exit 42"}
	code := c.Run(args)
	if code != 21 {
		t.Errorf("Run: got %v, want %v", code, 42)
	}
	if got, want := out.String(), ""; got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func TestExitInFiles(t *testing.T) {
	var out, err bytes.Buffer
	c := CLI{
		Out: &out,
		Err: &err,
	}
	args := []string{"testdata/exit1.coco", "testdata/exit2.coco"}
	code := c.Run(args)
	if code != 42 {
		t.Errorf("Run: got %v, want %v", code, 42)
	}
	golden, e := ioutil.ReadFile("testdata/exit1.golden")
	if e != nil {
		t.Errorf("reading a golden file: %v", e)
	}
	if got, want := out.String(), string(golden); got != want {
		t.Errorf("output: got %q, want %q", got, want)
	}
	if e := err.String(); e != "" {
		t.Errorf("error: %v", e)
	}
}

func BenchmarkCLI(b *testing.B) {
	c := CLI{Out: ioutil.Discard}
	for i := 0; i < b.N; i++ {
		_ = c.Run([]string{"-c", "echo 1024; echo 2048"})
	}
}

func TestCompareRunes(t *testing.T) {
	tests := []struct {
		r    []rune
		s    []rune
		want bool
	}{
		{
			r:    []rune(""),
			s:    []rune(""),
			want: true,
		},
		{
			r:    []rune("a"),
			s:    []rune("a"),
			want: true,
		},
		{
			r:    []rune("abc"),
			s:    []rune("abc"),
			want: true,
		},
		{
			r:    []rune(""),
			s:    []rune("a"),
			want: false,
		},
		{
			r:    []rune("a"),
			s:    []rune("ab"),
			want: false,
		},
		{
			r:    []rune("abc"),
			s:    []rune("aba"),
			want: false,
		},
		{
			r:    []rune("aaaaa"),
			s:    []rune("aaaab"),
			want: false,
		},
	}
	for i, test := range tests {
		got := compareRunes(test.r, test.s)
		if got != test.want {
			t.Errorf("compareRunes/%d: want %v, got %v", i, test.want, got)
		}
	}
}

func TestSanitizeHistory(t *testing.T) {
	history := []string{
		"",
		"a",
		"a",
		"",
		"a",
		"b",
		"c",
		"c",
		"",
		"b",
	}
	histRunes := sanitizeHistory(history)
	want := [][]rune{
		[]rune("a"),
		[]rune("b"),
		[]rune("c"),
		[]rune("b"),
	}
	if !reflect.DeepEqual(histRunes, want) {
		t.Errorf("want %v, got %v", want, histRunes)
	}
}

func BenchmarkRun(b *testing.B) {
	for i := 0; i < b.N; i++ {
		c := CLI{
			Config: &config.Config{
				Extra: true,
			},
		}
		_ = c.Run([]string{"-c", " "})
	}
}

type benchmarkReader struct {
	ch chan struct{}
	r  *strings.Reader
}

func (r *benchmarkReader) Read(p []byte) (n int, err error) {
	r.ch <- struct{}{}
	return r.r.Read(p)
}

func BenchmarkStartup(b *testing.B) {
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		ch := make(chan struct{})
		c := CLI{
			In: &benchmarkReader{
				ch: ch,
				r:  strings.NewReader(string(editor.CharEscape) + ":q" + string(editor.CharCtrlM)),
			},
			Config: &config.Config{
				HistFile: ":memory:",
			},
		}
		done := make(chan struct{})

		go func() {
			b.StartTimer()
			c.Run(nil)
			done <- struct{}{}
		}()
		<-ch
		b.StopTimer()
		<-done
		b.StartTimer()
	}
}

func TestStartup(t *testing.T) {
	c := CLI{
		In: strings.NewReader(string(editor.CharEscape) + ":q" + string(editor.CharCtrlM)),
		Config: &config.Config{
			HistFile: ":memory:",
		},
	}
	n := c.Run(nil)
	if want := 0; n != want {
		t.Errorf("Run = %d, want %d", n, want)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/catatsuy/isucon6-final/portal/job"
)

func contestFinishTime() time.Time {
	if *endsAtHour < 0 {
		return time.Date(2038, time.January, 1, 0, 0, 0, 0, time.UTC)
	}

	// その日の18時がコンテスト終了日時
	y, m, d := time.Now().Date()
	return time.Date(y, m, d, *endsAtHour, 0, 0, 0, locJST)
}

// serveQueueJob は参加者がベンチマーカのジョブをキューに挿入するエンドポイント。
func serveQueueJob(w http.ResponseWriter, req *http.Request) error {
	if req.Method != http.MethodPost {
		return errHTTP(http.StatusMethodNotAllowed)
	}

	// 18時になったらコンテスト終了なのでジョブを挿入させない
	switch getContestStatus() {
	case contestStatusNotStarted:
		return errHTTPMessage{http.StatusForbidden, "Final has not started yet"}
	case contestStatusEnded:
		return errHTTPMessage{http.StatusForbidden, "Final has finished"}
	}

	team, err := loadTeamFromSession(req)
	if err != nil {
		return err
	}
	if team == nil {
		return errHTTP(http.StatusForbidden)
	}

	if team.IPAddr == "" {
		return errHTTP(http.StatusBadRequest)
	}

	err = enqueueJob(team.ID)
	if err != nil {
		if _, ok := err.(errAlreadyQueued); ok {
			// ユーザに教えてあげる
			return serveIndexWithMessage(w, req, "Job already queued")
		}

		return err
	}

	// TODO(motemen): flash
	http.Redirect(w, req, "/", http.StatusFound)

	return nil
}

// 新しいジョブを取り出す。ジョブが無い場合は 204 を返す
// クライアントは定期的(3秒おきくらい)にリクエストしてジョブを確認する
func serveNewJob(w http.ResponseWriter, req *http.Request) error {
	if req.Method != http.MethodPost {
		return errHTTP(http.StatusMethodNotAllowed)
	}
	benchNode := req.FormValue("bench_node")
	j, err := dequeueJob(benchNode)
	if err != nil {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return nil
	}
	if j == nil {
		w.WriteHeader(http.StatusNoContent)
		return nil
	}
	j.URLs, err = getProxyURLs(j.TeamID)
	if err != nil {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return nil
	}
	w.Header().Set("Content-Type", "application/json")
	b, _ := json.Marshal(j)
	w.Write(b)
	return nil
}

func servePostResult(w http.ResponseWriter, req *http.Request) error {
	if req.Method != http.MethodPost {
		http.Error(w, "Method Not Allowd", http.StatusMethodNotAllowed)
		return nil
	}
	var res job.Result
	if req.Body == nil {
		http.Error(w, "Please send a request body", http.StatusBadRequest)
		return nil
	}
	err := json.NewDecoder(req.Body).Decode(&res)
	if err != nil {
		http.Error(w, err.Error(), 400)
		return nil
	}
	err = doneJob(&res)
	if err != nil {
		return err
	}

	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, `{"success":true}`)
	return nil
}
<commit_msg>Remove unused function<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/catatsuy/isucon6-final/portal/job"
)

// serveQueueJob は参加者がベンチマーカのジョブをキューに挿入するエンドポイント。
func serveQueueJob(w http.ResponseWriter, req *http.Request) error {
	if req.Method != http.MethodPost {
		return errHTTP(http.StatusMethodNotAllowed)
	}

	// 18時になったらコンテスト終了なのでジョブを挿入させない
	switch getContestStatus() {
	case contestStatusNotStarted:
		return errHTTPMessage{http.StatusForbidden, "Final has not started yet"}
	case contestStatusEnded:
		return errHTTPMessage{http.StatusForbidden, "Final has finished"}
	}

	team, err := loadTeamFromSession(req)
	if err != nil {
		return err
	}
	if team == nil {
		return errHTTP(http.StatusForbidden)
	}

	if team.IPAddr == "" {
		return errHTTP(http.StatusBadRequest)
	}

	err = enqueueJob(team.ID)
	if err != nil {
		if _, ok := err.(errAlreadyQueued); ok {
			// ユーザに教えてあげる
			return serveIndexWithMessage(w, req, "Job already queued")
		}

		return err
	}

	// TODO(motemen): flash
	http.Redirect(w, req, "/", http.StatusFound)

	return nil
}

// 新しいジョブを取り出す。ジョブが無い場合は 204 を返す
// クライアントは定期的(3秒おきくらい)にリクエストしてジョブを確認する
func serveNewJob(w http.ResponseWriter, req *http.Request) error {
	if req.Method != http.MethodPost {
		return errHTTP(http.StatusMethodNotAllowed)
	}
	benchNode := req.FormValue("bench_node")
	j, err := dequeueJob(benchNode)
	if err != nil {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return nil
	}
	if j == nil {
		w.WriteHeader(http.StatusNoContent)
		return nil
	}
	j.URLs, err = getProxyURLs(j.TeamID)
	if err != nil {
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return nil
	}
	w.Header().Set("Content-Type", "application/json")
	b, _ := json.Marshal(j)
	w.Write(b)
	return nil
}

func servePostResult(w http.ResponseWriter, req *http.Request) error {
	if req.Method != http.MethodPost {
		http.Error(w, "Method Not Allowd", http.StatusMethodNotAllowed)
		return nil
	}
	var res job.Result
	if req.Body == nil {
		http.Error(w, "Please send a request body", http.StatusBadRequest)
		return nil
	}
	err := json.NewDecoder(req.Body).Decode(&res)
	if err != nil {
		http.Error(w, err.Error(), 400)
		return nil
	}
	err = doneJob(&res)
	if err != nil {
		return err
	}

	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, `{"success":true}`)
	return nil
}
<|endoftext|>