<commit_before>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)
}
<commit_msg>config_test: fix error handling<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)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	dataDir := t.TempDir()
	cfgString := fmt.Sprintf(`
server:
  host_keys: [%v]
  tcpip_services:
    8080: HTTP
`, keyFile)
	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 cc

import (
	"flag"
	"os"
	"testing"
	"time"

	"github.com/damnever/cc/assert"
)

func TestConfigNew(t *testing.T) {
	{
		c := NewConfigFrom(map[string]interface{}{"foo": "bar"})
		assert.Check(t, c.Has("foo"), true)
	}
	{
		c, err := NewConfigFromJSON([]byte(`{"foo": 123}`))
		assert.Must(t, err)
		assert.Check(t, c.Has("foo"), true)
	}
	{
		c, err := NewConfigFromYAML([]byte(`name: good`))
		assert.Must(t, err)
		assert.Check(t, c.Has("name"), true)
	}
	{
		c, err := NewConfigFromFile("./example/example.yaml")
		assert.Must(t, err)
		assert.Check(t, c.Has("name"), true)
		assert.Check(t, c.Config("map").Has("key_one"), true)
		assert.Check(t, len(c.Value("map").Map()), 3)
		assert.Check(t, len(c.Value("list").List()), 4)
	}
	{
		c, err := NewConfigFromFile("./example/example.json")
		assert.Must(t, err)
		assert.Check(t, c.Has("name"), true)
		assert.Check(t, c.Config("map").Has("key_one"), true)
		assert.Check(t, len(c.Value("map").Map()), 3)
		assert.Check(t, len(c.Value("list").List()), 4)
	}
	if _, err := NewConfigFromFile("example/main.go"); err == nil {
		t.Fatal("expected error, got nothing")
	}
}

func TestConfigBasics(t *testing.T) {
	c := NewConfig()
	assert.Check(t, c.Has("foo"), false)
	c.SetDefault("foo", "bar")
	assert.Check(t, c.Has("foo"), true)
	c.Set("foo", "baz")
	assert.Check(t, c.String("foo"), "baz")
	c.SetDefault("bar", "foo")
	assert.Check(t, c.Has("bar"), true)

	ca := NewConfig()
	assert.Check(t, ca.Has("foo"), false)
	assert.Check(t, ca.Has("baz"), false)
	ca.Set("bar", "baz")
	assert.Check(t, ca.Has("bar"), true)
	assert.Check(t, ca.String("bar"), "baz")
	ca.Merge(c)
	assert.Check(t, ca.Has("foo"), true)
	assert.Check(t, ca.String("bar"), "foo")
}

func TestConfigMust(t *testing.T) {
	defer func() {
		if err := recover(); err == nil {
			t.Fatal("expect error, got nothing")
		}
	}()
	c := NewConfig()
	c.Must("must")
}

func TestConfigValue(t *testing.T) {
	c := NewConfig()
	if x, ok := c.Value("xx").(Valuer); !ok {
		t.Fatalf("expected Valuer, got %#v\n", x)
	}
}

func TestConfigPattern(t *testing.T) {
	c := NewConfig()
	if x, ok := c.Pattern("xx").(Patterner); !ok {
		t.Fatalf("expected Patterner, got %#v\n", x)
	}
}

func TestConfigGetConfig(t *testing.T) {
	c := NewConfig()
	c.Set("unknwn_map", map[interface{}]interface{}{"foo": "bar"})
	c.Set("string_map", map[string]interface{}{"foo": "bar"})
	cc := NewConfig()
	cc.Set("foo", "bar")
	c.Set("cc", cc)

	assert.Check(t, c.Has("unknwn_map"), true)
	assert.Check(t, c.Has("string_map"), true)
	assert.Check(t, c.Has("cc"), true)
	assert.Check(t, c.Config("unknwn_map").Has("foo"), true)
	assert.Check(t, c.Config("string_map").Has("foo"), true)
	assert.Check(t, c.Config("cc").Has("foo"), true)

	ccc := c.Config("non")
	ccc.Set("test", "good")
	ccc.Set("good", "bad")
	ccc = c.Config("non")
	assert.Check(t, len(c.Value("non").Map()), 2)
}

func TestConfigGetString(t *testing.T) {
	c := NewConfig()
	c.Set("string", "foo")
	assert.Check(t, c.Has("string"), true)
	assert.Check(t, c.String("string"), "foo")
	assert.Check(t, c.StringOr("string", "bar"), "foo")
	assert.Check(t, c.Has("foo"), false)
	assert.Check(t, c.String("foo"), "")
	assert.Check(t, c.StringOr("foo", "bar"), "bar")

	c.Set("www", "mmm")
	res, ok := c.StringAnd("www", "^m")
	assert.Check(t, ok, true)
	assert.Check(t, res, "mmm")
	res, ok = c.StringAnd("www", "^w")
	assert.Check(t, ok, false)
	assert.Check(t, res, "")
	res, ok = c.StringAnd("mmm", "^w")
	assert.Check(t, ok, false)
	assert.Check(t, res, "")

	assert.Check(t, c.StringAndOr("www", "^m", "lll"), "mmm")
	assert.Check(t, c.StringAndOr("www", "^w", "lll"), "lll")
	assert.Check(t, c.StringAndOr("mmm", "^w", "lll"), "lll")

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "string")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.StringOr("test_env", "XXX"), "string")

	flag.String("str_flag", "do", "usage")
	c.ParseFlags()
	assert.Check(t, c.String("str_flag"), "do")
}

func TestConfigGetBool(t *testing.T) {
	c := NewConfig()
	c.Set("bool", true)
	assert.Check(t, c.Has("bool"), true)
	assert.Check(t, c.Bool("bool"), true)
	assert.Check(t, c.BoolOr("bool", false), true)
	assert.Check(t, c.Has("non"), false)
	assert.Check(t, c.Bool("non"), false)
	assert.Check(t, c.BoolOr("non", true), true)

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "1")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.BoolOr("test_env", false), true)

	flag.Bool("bool_flag", true, "usage")
	c.ParseFlags()
	assert.Check(t, c.Bool("bool_flag"), true)
}

func TestConfigGetInt(t *testing.T) {
	c := NewConfig()
	c.Set("int", 33)
	assert.Check(t, c.Has("int"), true)
	assert.Check(t, c.Int("int"), 33)
	assert.Check(t, c.IntOr("int", 333), 33)
	res, ok := c.IntAnd("int", "N>3")
	assert.Check(t, ok, true)
	assert.Check(t, res, 33)

	assert.Check(t, c.Has("non"), false)
	assert.Check(t, c.Int("non"), 0)
	assert.Check(t, c.IntOr("non", 333), 333)
	res, ok = c.IntAnd("non", "N>3")
	assert.Check(t, ok, false)
	assert.Check(t, res, 0)

	assert.Check(t, c.IntAndOr("int", "N>3", 333), 33)
	assert.Check(t, c.IntAndOr("int", "N>33", 333), 333)
	assert.Check(t, c.IntAndOr("non", "N>33", 3333), 3333)

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "1111")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.IntOr("test_env", 11), 1111)

	flag.Int("int_flag", 32, "usage")
	flag.Int64("int64_flag", 64, "usage")
	c.ParseFlags()
	assert.Check(t, c.Int("int_flag"), 32)
	assert.Check(t, c.Int64("int64_flag"), int64(64))
}

func TestConfigGetFloat(t *testing.T) {
	c := NewConfig()
	c.Set("float", 333.3)
	assert.Check(t, c.Has("float"), true)
	assert.Check(t, c.Float("float"), 333.3)
	assert.Check(t, c.FloatOr("float", 33.33), 333.3)
	res, ok := c.FloatAnd("float", "N*10==3333")
	assert.Check(t, ok, true)
	assert.Check(t, res, 333.3)

	assert.Check(t, c.Has("non"), false)
	assert.Check(t, c.Float("non"), 0.0)
	assert.Check(t, c.FloatOr("non", 33.33), 33.33)
	res, ok = c.FloatAnd("non", "N>0")
	assert.Check(t, ok, false)
	assert.Check(t, res, 0.0)

	assert.Check(t, c.FloatAndOr("float", "N>33.3", 33.3), 333.3)
	assert.Check(t, c.FloatAndOr("float", "N>333.3", 33.3), 33.3)
	assert.Check(t, c.FloatAndOr("non", "N>33.3", 33.33), 33.33)

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "11.11")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.FloatOr("test_env", 1.1), 11.11)

	flag.Float64("float_flag", 64.64, "usage")
	c.ParseFlags()
	assert.Check(t, c.Float("float_flag"), 64.64)
}

func TestConfigGetDuration(t *testing.T) {
	c := NewConfig()
	c.Set("t", 300)
	assert.Check(t, c.Has("t"), true)
	assert.Check(t, c.Duration("t"), time.Duration(300))
	assert.Check(t, c.DurationOr("t", 333), time.Duration(300))
	assert.Check(t, c.Has("tt"), false)
	assert.Check(t, c.Duration("tt"), time.Duration(0))
	assert.Check(t, c.DurationOr("tt", 333), time.Duration(333))

	res, ok := c.DurationAnd("t", "N>30")
	assert.Check(t, ok, true)
	assert.Check(t, res, time.Duration(300))
	res, ok = c.DurationAnd("non", "N>1")
	assert.Check(t, ok, false)
	assert.Check(t, res, time.Duration(0))

	assert.Check(t, c.DurationAndOr("t", "N>30", 333), time.Duration(300))
	assert.Check(t, c.DurationAndOr("t", "N>300", 333), time.Duration(333))
	assert.Check(t, c.DurationAndOr("non", "N>33", 33), time.Duration(33))

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "1111")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.DurationOr("test_env", 11), time.Duration(1111))

	flag.Duration("duration_flag", 6464, "usage")
	c.ParseFlags()
	assert.Check(t, c.Duration("duration_flag"), time.Duration(6464))
}
<commit_msg>Existence check<commit_after>package cc

import (
	"flag"
	"os"
	"testing"
	"time"

	"github.com/damnever/cc/assert"
)

func TestConfigNew(t *testing.T) {
	{
		c := NewConfigFrom(map[string]interface{}{"foo": "bar"})
		assert.Check(t, c.Has("foo"), true)
	}
	{
		c, err := NewConfigFromJSON([]byte(`{"foo": 123}`))
		assert.Must(t, err)
		assert.Check(t, c.Has("foo"), true)
	}
	{
		c, err := NewConfigFromYAML([]byte(`name: good`))
		assert.Must(t, err)
		assert.Check(t, c.Has("name"), true)
	}
	{
		c, err := NewConfigFromFile("./example/example.yaml")
		assert.Must(t, err)
		assert.Check(t, c.Has("name"), true)
		assert.Check(t, c.Config("map").Has("key_one"), true)
		assert.Check(t, len(c.Value("map").Map()), 3)
		assert.Check(t, len(c.Value("list").List()), 4)
	}
	{
		c, err := NewConfigFromFile("./example/example.json")
		assert.Must(t, err)
		assert.Check(t, c.Has("name"), true)
		assert.Check(t, c.Config("map").Has("key_one"), true)
		assert.Check(t, len(c.Value("map").Map()), 3)
		assert.Check(t, len(c.Value("list").List()), 4)
	}
	if _, err := NewConfigFromFile("example/main.go"); err == nil {
		t.Fatal("expected error, got nothing")
	}
}

func TestConfigBasics(t *testing.T) {
	c := NewConfig()
	assert.Check(t, c.Has("foo"), false)
	c.SetDefault("foo", "bar")
	assert.Check(t, c.Has("foo"), true)
	c.Set("foo", "baz")
	assert.Check(t, c.String("foo"), "baz")
	c.SetDefault("bar", "foo")
	assert.Check(t, c.Has("bar"), true)

	ca := NewConfig()
	assert.Check(t, ca.Has("foo"), false)
	assert.Check(t, ca.Has("baz"), false)
	ca.Set("bar", "baz")
	assert.Check(t, ca.Has("bar"), true)
	assert.Check(t, ca.String("bar"), "baz")
	ca.Merge(c)
	assert.Check(t, ca.Has("foo"), true)
	assert.Check(t, ca.String("bar"), "foo")
}

func TestConfigMust(t *testing.T) {
	defer func() {
		if err := recover(); err == nil {
			t.Fatal("expect error, got nothing")
		}
	}()
	c := NewConfig()
	c.Must("must")
}

func TestConfigValue(t *testing.T) {
	c := NewConfig()
	if x, ok := c.Value("xx").(Valuer); !ok {
		t.Fatalf("expected Valuer, got %#v\n", x)
	}
}

func TestConfigPattern(t *testing.T) {
	c := NewConfig()
	if x, ok := c.Pattern("xx").(Patterner); !ok {
		t.Fatalf("expected Patterner, got %#v\n", x)
	}
}

func TestConfigGetConfig(t *testing.T) {
	c := NewConfig()
	c.Set("unknwn_map", map[interface{}]interface{}{"foo": "bar"})
	c.Set("string_map", map[string]interface{}{"foo": "bar"})
	cc := NewConfig()
	cc.Set("foo", "bar")
	c.Set("cc", cc)

	assert.Check(t, c.Has("unknwn_map"), true)
	assert.Check(t, c.Has("string_map"), true)
	assert.Check(t, c.Has("cc"), true)
	assert.Check(t, c.Config("unknwn_map").Has("foo"), true)
	assert.Check(t, c.Config("string_map").Has("foo"), true)
	assert.Check(t, c.Config("cc").Has("foo"), true)

	ccc := c.Config("non")
	ccc.Set("test", "good")
	ccc.Set("good", "bad")
	ccc = c.Config("non")
	assert.Check(t, len(c.Value("non").Map()), 2)
}

func TestConfigGetString(t *testing.T) {
	c := NewConfig()
	c.Set("string", "foo")
	assert.Check(t, c.Has("string"), true)
	assert.Check(t, c.String("string"), "foo")
	assert.Check(t, c.StringOr("string", "bar"), "foo")
	assert.Check(t, c.Has("foo"), false)
	assert.Check(t, c.String("foo"), "")
	assert.Check(t, c.StringOr("foo", "bar"), "bar")

	c.Set("www", "mmm")
	res, ok := c.StringAnd("www", "^m")
	assert.Check(t, ok, true)
	assert.Check(t, res, "mmm")
	res, ok = c.StringAnd("www", "^w")
	assert.Check(t, ok, false)
	assert.Check(t, res, "")
	res, ok = c.StringAnd("mmm", "^w")
	assert.Check(t, ok, false)
	assert.Check(t, res, "")

	assert.Check(t, c.StringAndOr("www", "^m", "lll"), "mmm")
	assert.Check(t, c.StringAndOr("www", "^w", "lll"), "lll")
	assert.Check(t, c.StringAndOr("mmm", "^w", "lll"), "lll")

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "string")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.StringOr("test_env", "XXX"), "string")

	flag.String("str_flag", "do", "usage")
	c.ParseFlags()
	assert.Check(t, c.Has("str_flag"), true)
	assert.Check(t, c.String("str_flag"), "do")
}

func TestConfigGetBool(t *testing.T) {
	c := NewConfig()
	c.Set("bool", true)
	assert.Check(t, c.Has("bool"), true)
	assert.Check(t, c.Bool("bool"), true)
	assert.Check(t, c.BoolOr("bool", false), true)
	assert.Check(t, c.Has("non"), false)
	assert.Check(t, c.Bool("non"), false)
	assert.Check(t, c.BoolOr("non", true), true)

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "1")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.BoolOr("test_env", false), true)

	flag.Bool("bool_flag", true, "usage")
	c.ParseFlags()
	assert.Check(t, c.Has("bool_flag"), true)
	assert.Check(t, c.Bool("bool_flag"), true)
}

func TestConfigGetInt(t *testing.T) {
	c := NewConfig()
	c.Set("int", 33)
	assert.Check(t, c.Has("int"), true)
	assert.Check(t, c.Int("int"), 33)
	assert.Check(t, c.IntOr("int", 333), 33)
	res, ok := c.IntAnd("int", "N>3")
	assert.Check(t, ok, true)
	assert.Check(t, res, 33)

	assert.Check(t, c.Has("non"), false)
	assert.Check(t, c.Int("non"), 0)
	assert.Check(t, c.IntOr("non", 333), 333)
	res, ok = c.IntAnd("non", "N>3")
	assert.Check(t, ok, false)
	assert.Check(t, res, 0)

	assert.Check(t, c.IntAndOr("int", "N>3", 333), 33)
	assert.Check(t, c.IntAndOr("int", "N>33", 333), 333)
	assert.Check(t, c.IntAndOr("non", "N>33", 3333), 3333)

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "1111")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.IntOr("test_env", 11), 1111)

	flag.Int("int_flag", 32, "usage")
	flag.Int64("int64_flag", 64, "usage")
	c.ParseFlags()
	assert.Check(t, c.Has("int_flag"), true)
	assert.Check(t, c.Int("int_flag"), 32)
	assert.Check(t, c.Has("int64_flag"), true)
	assert.Check(t, c.Int64("int64_flag"), int64(64))
}

func TestConfigGetFloat(t *testing.T) {
	c := NewConfig()
	c.Set("float", 333.3)
	assert.Check(t, c.Has("float"), true)
	assert.Check(t, c.Float("float"), 333.3)
	assert.Check(t, c.FloatOr("float", 33.33), 333.3)
	res, ok := c.FloatAnd("float", "N*10==3333")
	assert.Check(t, ok, true)
	assert.Check(t, res, 333.3)

	assert.Check(t, c.Has("non"), false)
	assert.Check(t, c.Float("non"), 0.0)
	assert.Check(t, c.FloatOr("non", 33.33), 33.33)
	res, ok = c.FloatAnd("non", "N>0")
	assert.Check(t, ok, false)
	assert.Check(t, res, 0.0)

	assert.Check(t, c.FloatAndOr("float", "N>33.3", 33.3), 333.3)
	assert.Check(t, c.FloatAndOr("float", "N>333.3", 33.3), 33.3)
	assert.Check(t, c.FloatAndOr("non", "N>33.3", 33.33), 33.33)

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "11.11")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.FloatOr("test_env", 1.1), 11.11)

	flag.Float64("float_flag", 64.64, "usage")
	c.ParseFlags()
	assert.Check(t, c.Has("float_flag"), true)
	assert.Check(t, c.Float("float_flag"), 64.64)
}

func TestConfigGetDuration(t *testing.T) {
	c := NewConfig()
	c.Set("t", 300)
	assert.Check(t, c.Has("t"), true)
	assert.Check(t, c.Duration("t"), time.Duration(300))
	assert.Check(t, c.DurationOr("t", 333), time.Duration(300))
	assert.Check(t, c.Has("tt"), false)
	assert.Check(t, c.Duration("tt"), time.Duration(0))
	assert.Check(t, c.DurationOr("tt", 333), time.Duration(333))

	res, ok := c.DurationAnd("t", "N>30")
	assert.Check(t, ok, true)
	assert.Check(t, res, time.Duration(300))
	res, ok = c.DurationAnd("non", "N>1")
	assert.Check(t, ok, false)
	assert.Check(t, res, time.Duration(0))

	assert.Check(t, c.DurationAndOr("t", "N>30", 333), time.Duration(300))
	assert.Check(t, c.DurationAndOr("t", "N>300", 333), time.Duration(333))
	assert.Check(t, c.DurationAndOr("non", "N>33", 33), time.Duration(33))

	assert.Check(t, c.Has("test_env"), false)
	os.Setenv("test_env", "1111")
	defer func() { os.Unsetenv("test_env") }()
	assert.Check(t, c.Has("test_env"), true)
	assert.Check(t, c.DurationOr("test_env", 11), time.Duration(1111))

	flag.Duration("duration_flag", 6464, "usage")
	c.ParseFlags()
	assert.Check(t, c.Has("duration_flag"), true)
	assert.Check(t, c.Duration("duration_flag"), time.Duration(6464))
}
<|endoftext|>


<commit_before>// Package oplogc provides an easy to use client interface for the oplog service.
//
// See https://github.com/dailymotion/oplog for more information on oplog.
//
// In case of a connection failure recovery the ack mechanism allows you to handle operations in parallel
// without loosing track of which operation has been handled.
//
// See cmd/oplog-tail for another usage example.
package oplogc

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"
)

// Options is the subscription options
type Options struct {
	// Path of the state file where to persiste the current oplog position.
	// If empty string, the state is not stored.
	StateFile string
	// AllowReplication activates replication if the state file is not found.
	// When false, a consumer with no state file will only get future operations.
	AllowReplication bool
	// Password to access password protected oplog
	Password string
	// Filters to apply on the oplog output
	Filter Filter
}

// Filter contains arguments to filter the oplog output
type Filter struct {
	// A list of types to filter on
	Types []string
	// A list of parent type/id to filter on
	Parents []string
}

// Consumer holds all the information required to connect to an oplog server
type Consumer struct {
	// URL of the oplog
	url string
	// options for the consumer's subscription
	options Options
	// lastID is the current most advanced acked event id
	lastID string
	// saved is true when current lastID is persisted
	saved bool
	// processing is true when a process loop is in progress
	processing bool
	// mu is a mutex used to coordinate access to lastID and saved properties
	mu *sync.RWMutex
	// http is the client used to connect to the oplog
	http http.Client
	// body points to the current streamed response body
	body io.ReadCloser
	// ife holds all event ids sent to the consumer but no yet acked
	ife *inFlightEvents
	// ack is a channel to ack the operations
	ack chan Operation
	// stop is a channel used to stop the process loop
	stop chan struct{}
}

// ErrAccessDenied is returned by Subscribe when the oplog requires a password
// different from the one provided in options.
var ErrAccessDenied = errors.New("invalid credentials")

// ErrResumeFailed is returned when the requested last id was not found by the
// oplog server. This may happen when the last id is very old or size of the
// oplog capped collection is too small for the load.
//
// When this error happen, the consumer may choose to either ignore the lost events
// or force a full replication.
var ErrResumeFailed = errors.New("resume failed")

// ErrorWritingState is returned when the last processed id can't be written to
// the state file.
var ErrWritingState = errors.New("writing state file failed")

// Subscribe creates a Consumer to connect to the given URL.
func Subscribe(url string, options Options) *Consumer {
	qs := ""
	if len(options.Filter.Parents) > 0 {
		parents := strings.Join(options.Filter.Parents, ",")
		if parents != "" {
			qs += "?parents="
			qs += parents
		}
	}
	if len(options.Filter.Types) > 0 {
		types := strings.Join(options.Filter.Types, ",")
		if types != "" {
			if qs == "" {
				qs += "?"
			} else {
				qs += "&"
			}
			qs += "types="
			qs += types
		}
	}

	c := &Consumer{
		url:     strings.Join([]string{url, qs}, ""),
		options: options,
		ife:     newInFlightEvents(),
		mu:      &sync.RWMutex{},
		ack:     make(chan Operation),
	}

	return c
}

// Start reads the oplog output and send operations back thru the returned ops channel.
// The caller must then call the Done() method on operation when it has been handled.
// Failing to call Done() the operations would prevent any resume in case of connection
// failure or restart of the process.
//
// Any errors are return on the errs channel. In all cases, the Start() method will
// try to reconnect and/or ignore the error. It is the callers responsability to stop
// the process loop by calling the Stop() method.
//
// When the loop has ended, a message is sent thru the done channel.
func (c *Consumer) Start() (ops chan Operation, errs chan error, done chan bool) {
	ops = make(chan Operation)
	errs = make(chan error)
	done = make(chan bool)

	// Ensure we never have more than one process loop running
	if c.processing {
		panic("Can't run two process loops in parallel")
	}
	c.processing = true

	c.mu.Lock()
	c.stop = make(chan struct{})
	stop := c.stop
	c.mu.Unlock()

	// Recover the last event id saved from a previous excution
	lastID, err := c.loadLastEventID()
	if err != nil {
		errs <- err
		return
	}
	c.lastID = lastID

	wg := sync.WaitGroup{}

	// SSE stream reading
	stopReadStream := make(chan struct{}, 1)
	wg.Add(1)
	go c.readStream(ops, errs, stopReadStream, &wg)

	// Periodic (non blocking) saving of the last id when needed
	stopStateSaving := make(chan struct{}, 1)
	if c.options.StateFile != "" {
		wg.Add(1)
		go c.periodicStateSaving(errs, stopStateSaving, &wg)
	}

	go func() {
		for {
			select {
			case <-stop:
				// If a stop is requested, we ensure all go routines are stopped
				close(stopReadStream)
				close(stopStateSaving)
				if c.body != nil {
					// Closing the body will ensure readStream isn't blocked in IO wait
					c.body.Close()
				}
				wg.Wait()
				c.processing = false
				done <- true
				return
			case op := <-c.ack:
				if op.Event == "reset" {
					c.ife.Unlock()
				}
				if idx := c.ife.pull(op.ID); idx == 0 {
					c.SetLastID(op.ID)
				}
			}
		}
	}()

	return
}

// Stop instructs the Start() loop to stop
func (c *Consumer) Stop() {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.stop != nil {
		close(c.stop)
		c.stop = nil
	}
}

// readStream maintains a connection to the oplog stream and read sent events as they are coming
func (c *Consumer) readStream(ops chan<- Operation, errs chan<- error, stop <-chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()

	c.connect()
	d := newDecoder(c.body)
	op := Operation{}
	op.ack = c.ack
	backoff := time.Second
	for {
		err := d.next(&op)
		select {
		case <-stop:
			return
		default:
			// proceed
		}
		if err != nil {
			errs <- err
			for {
				time.Sleep(backoff)
				if backoff < 30*time.Second {
					backoff *= 2
				}
				if err = c.connect(); err == nil {
					d = newDecoder(c.body)
					break
				}
				errs <- err
			}
			continue
		}

		c.ife.push(op.ID)
		if op.Event == "reset" {
			// We must not process any further operation until the "reset" operation
			// is not acke
			c.ife.Lock()
		}
		select {
		case <-stop:
			return
		default:
			ops <- op
		}

		// reset backoff on success
		backoff = time.Second
	}
}

// periodicStateSaving saves the lastID into a file every seconds if it has been updated
func (c *Consumer) periodicStateSaving(errs chan<- error, stop <-chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()

	for {
		select {
		case <-stop:
			return
		case <-time.After(time.Second):
			c.mu.RLock()
			saved := c.saved
			lastID := c.lastID
			c.mu.RUnlock()
			if saved {
				continue
			}
			if err := c.saveLastEventID(lastID); err != nil {
				errs <- ErrWritingState
			}
			c.mu.Lock()
			c.saved = lastID == c.lastID
			c.mu.Unlock()
		}
	}
}

// LastID returns the most advanced acked event id
func (c *Consumer) LastID() string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.lastID
}

// SetLastID sets the last id to the given value and informs the save go routine
func (c *Consumer) SetLastID(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.lastID = id
	c.saved = false
}

// connect tries to connect to the oplog event stream
func (c *Consumer) connect() (err error) {
	if c.body != nil {
		c.body.Close()
	}
	// Usable dummy body in case of connection error
	c.body = ioutil.NopCloser(bytes.NewBuffer([]byte{}))

	req, err := http.NewRequest("GET", c.url, nil)
	if err != nil {
		return
	}
	req.Header.Set("Cache-Control", "no-cache")
	req.Header.Set("Accept", "text/event-stream")
	lastID := c.LastID()
	if len(lastID) > 0 {
		req.Header.Set("Last-Event-ID", lastID)
	}
	if c.options.Password != "" {
		req.SetBasicAuth("", c.options.Password)
	}
	res, err := c.http.Do(req)
	if err != nil {
		return
	}
	if res.StatusCode == 403 || res.StatusCode == 401 {
		err = ErrAccessDenied
		return
	}
	if res.StatusCode != 200 {
		message, _ := ioutil.ReadAll(res.Body)
		err = fmt.Errorf("HTTP error %d: %s", res.StatusCode, string(message))
		return
	}
	c.body = res.Body
	return
}

// loadLastEventID tries to read the last event id from the state file.
//
// If the StateFile option was not set, the id will always be an empty string
// as for tailing only future events.
//
// If the StateFile option is set but no file exists, the last event id is
// initialized to "0" in order to request a full replication if AllowReplication
// option is set to true or to an empty string otherwise (start at present).
func (c *Consumer) loadLastEventID() (id string, err error) {
	if c.options.StateFile == "" {
		return "", nil
	}
	_, err = os.Stat(c.options.StateFile)
	if os.IsNotExist(err) {
		if c.options.AllowReplication {
			// full replication
			id = "0"
		} else {
			// start at NOW()
			id = ""
		}
		err = nil
	} else if err == nil {
		var content []byte
		content, err = ioutil.ReadFile(c.options.StateFile)
		if err != nil {
			return
		}
		if match, _ := regexp.Match("^(?:[0-9]{0,13}|[0-9a-f]{24})$", content); !match {
			err = errors.New("state file contains invalid data")
		}
		id = string(content)
	}
	return
}

// saveLastEventID persiste the last event id into a file
func (c *Consumer) saveLastEventID(id string) error {
	return ioutil.WriteFile(c.options.StateFile, []byte(id), 0644)
}
<commit_msg>Add proxy support<commit_after>// Package oplogc provides an easy to use client interface for the oplog service.
//
// See https://github.com/dailymotion/oplog for more information on oplog.
//
// In case of a connection failure recovery the ack mechanism allows you to handle operations in parallel
// without loosing track of which operation has been handled.
//
// See cmd/oplog-tail for another usage example.
package oplogc

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	neturl "net/url"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"
)

// Options is the subscription options
type Options struct {
	// Path of the state file where to persiste the current oplog position.
	// If empty string, the state is not stored.
	StateFile string
	// AllowReplication activates replication if the state file is not found.
	// When false, a consumer with no state file will only get future operations.
	AllowReplication bool
	// Password to access password protected oplog
	Password string
	// Proxy to be used to access oplog
	Proxy string
	// Filters to apply on the oplog output
	Filter Filter
}

// Filter contains arguments to filter the oplog output
type Filter struct {
	// A list of types to filter on
	Types []string
	// A list of parent type/id to filter on
	Parents []string
}

// Consumer holds all the information required to connect to an oplog server
type Consumer struct {
	// URL of the oplog
	url string
	// options for the consumer's subscription
	options Options
	// lastID is the current most advanced acked event id
	lastID string
	// saved is true when current lastID is persisted
	saved bool
	// processing is true when a process loop is in progress
	processing bool
	// mu is a mutex used to coordinate access to lastID and saved properties
	mu *sync.RWMutex
	// http is the client used to connect to the oplog
	http http.Client
	// body points to the current streamed response body
	body io.ReadCloser
	// ife holds all event ids sent to the consumer but no yet acked
	ife *inFlightEvents
	// ack is a channel to ack the operations
	ack chan Operation
	// stop is a channel used to stop the process loop
	stop chan struct{}
}

// ErrAccessDenied is returned by Subscribe when the oplog requires a password
// different from the one provided in options.
var ErrAccessDenied = errors.New("invalid credentials")

// ErrResumeFailed is returned when the requested last id was not found by the
// oplog server. This may happen when the last id is very old or size of the
// oplog capped collection is too small for the load.
//
// When this error happen, the consumer may choose to either ignore the lost events
// or force a full replication.
var ErrResumeFailed = errors.New("resume failed")

// ErrorWritingState is returned when the last processed id can't be written to
// the state file.
var ErrWritingState = errors.New("writing state file failed")

// Subscribe creates a Consumer to connect to the given URL.
func Subscribe(url string, options Options) *Consumer {
	qs := ""
	if len(options.Filter.Parents) > 0 {
		parents := strings.Join(options.Filter.Parents, ",")
		if parents != "" {
			qs += "?parents="
			qs += parents
		}
	}
	if len(options.Filter.Types) > 0 {
		types := strings.Join(options.Filter.Types, ",")
		if types != "" {
			if qs == "" {
				qs += "?"
			} else {
				qs += "&"
			}
			qs += "types="
			qs += types
		}
	}

	var proxyFunc func(*http.Request) (*neturl.URL, error) = nil

	if len(options.Proxy) > 0 {
		url_proxy, _ := neturl.ParseRequestURI(options.Proxy)
		proxyFunc = http.ProxyURL(url_proxy)
	}

	c := &Consumer{
		url:     strings.Join([]string{url, qs}, ""),
		options: options,
		ife:     newInFlightEvents(),
		mu:      &sync.RWMutex{},
		ack:     make(chan Operation),
		http:    http.Client{Transport: &http.Transport{Proxy: proxyFunc}},
	}

	return c
}

// Start reads the oplog output and send operations back thru the returned ops channel.
// The caller must then call the Done() method on operation when it has been handled.
// Failing to call Done() the operations would prevent any resume in case of connection
// failure or restart of the process.
//
// Any errors are return on the errs channel. In all cases, the Start() method will
// try to reconnect and/or ignore the error. It is the callers responsability to stop
// the process loop by calling the Stop() method.
//
// When the loop has ended, a message is sent thru the done channel.
func (c *Consumer) Start() (ops chan Operation, errs chan error, done chan bool) {
	ops = make(chan Operation)
	errs = make(chan error)
	done = make(chan bool)

	// Ensure we never have more than one process loop running
	if c.processing {
		panic("Can't run two process loops in parallel")
	}
	c.processing = true

	c.mu.Lock()
	c.stop = make(chan struct{})
	stop := c.stop
	c.mu.Unlock()

	// Recover the last event id saved from a previous excution
	lastID, err := c.loadLastEventID()
	if err != nil {
		errs <- err
		return
	}
	c.lastID = lastID

	wg := sync.WaitGroup{}

	// SSE stream reading
	stopReadStream := make(chan struct{}, 1)
	wg.Add(1)
	go c.readStream(ops, errs, stopReadStream, &wg)

	// Periodic (non blocking) saving of the last id when needed
	stopStateSaving := make(chan struct{}, 1)
	if c.options.StateFile != "" {
		wg.Add(1)
		go c.periodicStateSaving(errs, stopStateSaving, &wg)
	}

	go func() {
		for {
			select {
			case <-stop:
				// If a stop is requested, we ensure all go routines are stopped
				close(stopReadStream)
				close(stopStateSaving)
				if c.body != nil {
					// Closing the body will ensure readStream isn't blocked in IO wait
					c.body.Close()
				}
				wg.Wait()
				c.processing = false
				done <- true
				return
			case op := <-c.ack:
				if op.Event == "reset" {
					c.ife.Unlock()
				}
				if idx := c.ife.pull(op.ID); idx == 0 {
					c.SetLastID(op.ID)
				}
			}
		}
	}()

	return
}

// Stop instructs the Start() loop to stop
func (c *Consumer) Stop() {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.stop != nil {
		close(c.stop)
		c.stop = nil
	}
}

// readStream maintains a connection to the oplog stream and read sent events as they are coming
func (c *Consumer) readStream(ops chan<- Operation, errs chan<- error, stop <-chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()

	c.connect()
	d := newDecoder(c.body)
	op := Operation{}
	op.ack = c.ack
	backoff := time.Second
	for {
		err := d.next(&op)
		select {
		case <-stop:
			return
		default:
			// proceed
		}
		if err != nil {
			errs <- err
			for {
				time.Sleep(backoff)
				if backoff < 30*time.Second {
					backoff *= 2
				}
				if err = c.connect(); err == nil {
					d = newDecoder(c.body)
					break
				}
				errs <- err
			}
			continue
		}

		c.ife.push(op.ID)
		if op.Event == "reset" {
			// We must not process any further operation until the "reset" operation
			// is not acke
			c.ife.Lock()
		}
		select {
		case <-stop:
			return
		default:
			ops <- op
		}

		// reset backoff on success
		backoff = time.Second
	}
}

// periodicStateSaving saves the lastID into a file every seconds if it has been updated
func (c *Consumer) periodicStateSaving(errs chan<- error, stop <-chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()

	for {
		select {
		case <-stop:
			return
		case <-time.After(time.Second):
			c.mu.RLock()
			saved := c.saved
			lastID := c.lastID
			c.mu.RUnlock()
			if saved {
				continue
			}
			if err := c.saveLastEventID(lastID); err != nil {
				errs <- ErrWritingState
			}
			c.mu.Lock()
			c.saved = lastID == c.lastID
			c.mu.Unlock()
		}
	}
}

// LastID returns the most advanced acked event id
func (c *Consumer) LastID() string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.lastID
}

// SetLastID sets the last id to the given value and informs the save go routine
func (c *Consumer) SetLastID(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.lastID = id
	c.saved = false
}

// connect tries to connect to the oplog event stream
func (c *Consumer) connect() (err error) {
	if c.body != nil {
		c.body.Close()
	}
	// Usable dummy body in case of connection error
	c.body = ioutil.NopCloser(bytes.NewBuffer([]byte{}))

	req, err := http.NewRequest("GET", c.url, nil)
	if err != nil {
		return
	}
	req.Header.Set("Cache-Control", "no-cache")
	req.Header.Set("Accept", "text/event-stream")
	lastID := c.LastID()
	if len(lastID) > 0 {
		req.Header.Set("Last-Event-ID", lastID)
	}
	if c.options.Password != "" {
		req.SetBasicAuth("", c.options.Password)
	}
	res, err := c.http.Do(req)
	if err != nil {
		return
	}
	if res.StatusCode == 403 || res.StatusCode == 401 {
		err = ErrAccessDenied
		return
	}
	if res.StatusCode != 200 {
		message, _ := ioutil.ReadAll(res.Body)
		err = fmt.Errorf("HTTP error %d: %s", res.StatusCode, string(message))
		return
	}
	c.body = res.Body
	return
}

// loadLastEventID tries to read the last event id from the state file.
//
// If the StateFile option was not set, the id will always be an empty string
// as for tailing only future events.
//
// If the StateFile option is set but no file exists, the last event id is
// initialized to "0" in order to request a full replication if AllowReplication
// option is set to true or to an empty string otherwise (start at present).
func (c *Consumer) loadLastEventID() (id string, err error) {
	if c.options.StateFile == "" {
		return "", nil
	}
	_, err = os.Stat(c.options.StateFile)
	if os.IsNotExist(err) {
		if c.options.AllowReplication {
			// full replication
			id = "0"
		} else {
			// start at NOW()
			id = ""
		}
		err = nil
	} else if err == nil {
		var content []byte
		content, err = ioutil.ReadFile(c.options.StateFile)
		if err != nil {
			return
		}
		if match, _ := regexp.Match("^(?:[0-9]{0,13}|[0-9a-f]{24})$", content); !match {
			err = errors.New("state file contains invalid data")
		}
		id = string(content)
	}
	return
}

// saveLastEventID persiste the last event id into a file
func (c *Consumer) saveLastEventID(id string) error {
	return ioutil.WriteFile(c.options.StateFile, []byte(id), 0644)
}
<|endoftext|>


<commit_before>// tsadmin/database
package database

import (
	"database/sql"
	"fmt"
	"log"
	"strconv"

	_ "github.com/go-sql-driver/mysql"
)

type Database struct {
	Name     string `json:"name"`
	Host     string `json:"host"`
	Port     int    `json:"port"`
	User     string `json:"username"`
	Password string `json:"password"`
}

type DatabaseStatus struct {
	Metadata  DatabaseMetadata  `json:"metadata"`
	Metrics   DatabaseMetrics   `json:"metrics"`
	Variables DatabaseVariables `json:"variables"`
}

type DatabaseMetadata struct {
	Name string `json:"name"`
	Host string `json:"host"`
	Port int    `json:"port"`
}

type DatabaseMetrics struct {
	CurrentConnections          int `json:"current_connections"`
	Connections                 int `json:"connections"`
	ConnectionsPerSecond        int `json:"connections_per_second"`
	AbortedConnections          int `json:"aborted_connections"`
	AbortedConnectionsPerSecond int `json:"aborted_connections_per_second"`
	Queries                     int `json:"queries"`
	QueriesPerSecond            int `json:"queries_per_second"`
	Reads                       int `json:"reads"`
	ReadsPerSecond              int `json:"reads_per_second"`
	Writes                      int `json:"writes"`
	WritesPerSecond             int `json:"writes_per_second"`
	Uptime                      int `json:"uptime"`
}

type DatabaseVariables struct {
	MaxConnections int `json:"max_connections"`
}

func (db *Database) String() string {
	return fmt.Sprintf("%s:%s@tcp(%s:%d)/information_schema", db.User, db.Password, db.Host, db.Port)
}

func Status(db Database, previous *DatabaseStatus) (*DatabaseStatus, error) {
	status := &DatabaseStatus{
		Metadata: DatabaseMetadata{
			Name: db.Name,
			Host: db.Host,
			Port: db.Port,
		},
		Metrics:   DatabaseMetrics{},
		Variables: DatabaseVariables{},
	}

	// Fetch the metrics
	status, _ = execQuery(db, "metrics", previous, status)

	// Fetch the variables
	status, _ = execQuery(db, "variables", previous, status)

	return status, nil
}

func execQuery(db Database, queryType string, previous *DatabaseStatus, status *DatabaseStatus) (*DatabaseStatus, error) {
	var (
		key   string
		value string
		table string
	)

	// Fetch all the db metrics/variables
	if queryType == "metrics" {
		table = "GLOBAL_STATUS"
	} else if queryType == "variables" {
		table = "GLOBAL_VARIABLES"
	} else {
		log.Fatal("Unknown queryType")
	}

	// Connect to the database
	conn, _ := sql.Open("mysql", db.String())
	defer conn.Close()

	// Fetch all the db metrics
	rows, err := conn.Query(fmt.Sprintf("SELECT VARIABLE_NAME AS 'key', VARIABLE_VALUE AS 'value' FROM %s", table))

	// Handle query errors
	if err != nil {
		return status, err
	}

	defer rows.Close()

	// Loop each metric/variable in the server status
	for rows.Next() {
		err := rows.Scan(&key, &value)

		// Handle row reading errors
		if err != nil {
			return status, err
		}

		// Process the metrics/variables
		if queryType == "metrics" {
			status, _ = processMetrics(previous, status, key, value)
		} else {
			status, _ = processVariables(status, key, value)
		}
	}

	// Check for any remaining errors
	err = rows.Err()

	return status, err
}

// Process metrics returned from the GLOBAL_STATUS table
func processMetrics(previous *DatabaseStatus, status *DatabaseStatus, key string, value string) (*DatabaseStatus, error) {
	var (
		err                error
		currentConnections int
		connections        int
		diff               int
		abortedConnections int
		uptime             int
		queries            int
	)

	switch key {
	// Current connections
	case "THREADS_CONNECTED":
		currentConnections, err = strconv.Atoi(value)
		status.Metrics.CurrentConnections = currentConnections
	// Connections per second
	case "CONNECTIONS":
		connections, err = strconv.Atoi(value)

		// If we don't have a previous value for the total connections
		// then cps is technically 0 as we don't know it yet
		if previous == nil || previous.Metrics.Connections == 0 {
			status.Metrics.ConnectionsPerSecond = 0
			status.Metrics.Connections = connections
			// Otherwise the value of cps is the diff between the current
			// and previous count of connections
		} else {
			diff = connections - previous.Metrics.Connections

			// qps can never be below 0..
			if diff > 0 {
				status.Metrics.ConnectionsPerSecond = diff
			} else {
				status.Metrics.ConnectionsPerSecond = 0
			}

			status.Metrics.Connections = connections
		}
	// Aborted connections per second
	case "ABORTED_CONNECTS":
		abortedConnections, err = strconv.Atoi(value)

		// If we don't have a previous value for the total aborted connections
		// then acps is technically 0 as we don't know it yet
		if previous == nil || previous.Metrics.AbortedConnections == 0 {
			status.Metrics.AbortedConnectionsPerSecond = 0
			status.Metrics.AbortedConnections = abortedConnections
			// Otherwise the value of acps is the diff between the current
			// and previous count of connections
		} else {
			diff = abortedConnections - previous.Metrics.AbortedConnections

			// qps can never be below 0..
			if diff > 0 {
				status.Metrics.AbortedConnectionsPerSecond = diff
			} else {
				status.Metrics.AbortedConnectionsPerSecond = 0
			}

			status.Metrics.AbortedConnections = abortedConnections
		}
	// Uptime
	case "UPTIME":
		uptime, err = strconv.Atoi(value)
		status.Metrics.Uptime = uptime
	// Queries per second
	case "QUERIES":
		queries, err = strconv.Atoi(value)

		// If we don't have a previous value for the total queries
		// then qps is technically 0 as we don't know it yet
		if previous == nil || previous.Metrics.Queries == 0 {
			status.Metrics.QueriesPerSecond = 0
			status.Metrics.Queries = queries
			// Otherwise the value of qps is the diff between the current
			// and previous count of queries
		} else {
			diff = queries - previous.Metrics.Queries

			// qps can never be below 0..
			if diff > 0 {
				status.Metrics.QueriesPerSecond = diff
			} else {
				status.Metrics.QueriesPerSecond = 0
			}

			status.Metrics.Queries = queries
		}
	}

	if err != nil {
		return status, err
	} else {
		return status, nil
	}
}

// Process variables returned from the GLOBAL_VARIABLES table
func processVariables(status *DatabaseStatus, key string, value string) (*DatabaseStatus, error) {
	var (
		err            error
		maxConnections int
	)

	// Max allowed connections
	if key == "MAX_CONNECTIONS" {
		maxConnections, err = strconv.Atoi(value)
		status.Variables.MaxConnections = maxConnections
	}

	if err != nil {
		return status, err
	} else {
		return status, nil
	}
}
<commit_msg>Reads per second<commit_after>// tsadmin/database
package database

import (
	"database/sql"
	"fmt"
	"log"
	"strconv"

	_ "github.com/go-sql-driver/mysql"
)

type Database struct {
	Name     string `json:"name"`
	Host     string `json:"host"`
	Port     int    `json:"port"`
	User     string `json:"username"`
	Password string `json:"password"`
}

type DatabaseStatus struct {
	Metadata  DatabaseMetadata  `json:"metadata"`
	Metrics   DatabaseMetrics   `json:"metrics"`
	Variables DatabaseVariables `json:"variables"`
}

type DatabaseMetadata struct {
	Name string `json:"name"`
	Host string `json:"host"`
	Port int    `json:"port"`
}

type DatabaseMetrics struct {
	CurrentConnections          int `json:"current_connections"`
	Connections                 int `json:"connections"`
	ConnectionsPerSecond        int `json:"connections_per_second"`
	AbortedConnections          int `json:"aborted_connections"`
	AbortedConnectionsPerSecond int `json:"aborted_connections_per_second"`
	Queries                     int `json:"queries"`
	QueriesPerSecond            int `json:"queries_per_second"`
	Reads                       int `json:"reads"`
	ReadsPerSecond              int `json:"reads_per_second"`
	Writes                      int `json:"writes"`
	WritesPerSecond             int `json:"writes_per_second"`
	Uptime                      int `json:"uptime"`
}

type DatabaseVariables struct {
	MaxConnections int `json:"max_connections"`
}

func (db *Database) String() string {
	return fmt.Sprintf("%s:%s@tcp(%s:%d)/information_schema", db.User, db.Password, db.Host, db.Port)
}

func Status(db Database, previous *DatabaseStatus) (*DatabaseStatus, error) {
	status := &DatabaseStatus{
		Metadata: DatabaseMetadata{
			Name: db.Name,
			Host: db.Host,
			Port: db.Port,
		},
		Metrics:   DatabaseMetrics{},
		Variables: DatabaseVariables{},
	}

	// Fetch the metrics
	status, _ = execQuery(db, "metrics", previous, status)

	// Fetch the variables
	status, _ = execQuery(db, "variables", previous, status)

	return status, nil
}

// Execute a query on the given database for looking up metrics/variables
func execQuery(db Database, queryType string, previous *DatabaseStatus, status *DatabaseStatus) (*DatabaseStatus, error) {
	var (
		key   string
		value string
		table string
	)

	// Fetch all the db metrics/variables
	if queryType == "metrics" {
		table = "GLOBAL_STATUS"
	} else if queryType == "variables" {
		table = "GLOBAL_VARIABLES"
	} else {
		log.Fatal("Unknown queryType")
	}

	// Connect to the database
	conn, _ := sql.Open("mysql", db.String())
	defer conn.Close()

	// Fetch all the db metrics
	rows, err := conn.Query(fmt.Sprintf("SELECT VARIABLE_NAME AS 'key', VARIABLE_VALUE AS 'value' FROM %s", table))

	// Handle query errors
	if err != nil {
		return status, err
	}

	defer rows.Close()

	// Loop each metric/variable in the server status
	for rows.Next() {
		err := rows.Scan(&key, &value)

		// Handle row reading errors
		if err != nil {
			return status, err
		}

		// Process the metrics/variables
		if queryType == "metrics" {
			status, _ = processMetrics(previous, status, key, value)
		} else {
			status, _ = processVariables(status, key, value)
		}
	}

	// Check for any remaining errors
	err = rows.Err()

	return status, err
}

// Process metrics returned from the GLOBAL_STATUS table
func processMetrics(previous *DatabaseStatus, status *DatabaseStatus, key string, value string) (*DatabaseStatus, error) {
	var (
		err                error
		currentConnections int
		connections        int
		diff               int
		abortedConnections int
		queries            int
		reads              int
		uptime             int
	)

	switch key {
	// Current connections
	case "THREADS_CONNECTED":
		currentConnections, err = strconv.Atoi(value)
		status.Metrics.CurrentConnections = currentConnections
	// Connections per second
	case "CONNECTIONS":
		connections, err = strconv.Atoi(value)

		// If we don't have a previous value for the total connections
		// then cps is technically 0 as we don't know it yet
		if previous == nil || previous.Metrics.Connections == 0 {
			status.Metrics.ConnectionsPerSecond = 0
			status.Metrics.Connections = connections
		// Otherwise the value of cps is the diff between the current
		// and previous count of connections
		} else {
			diff = connections - previous.Metrics.Connections

			// qps can never be below 0..
			if diff > 0 {
				status.Metrics.ConnectionsPerSecond = diff
			} else {
				status.Metrics.ConnectionsPerSecond = 0
			}

			status.Metrics.Connections = connections
		}
	// Aborted connections per second
	case "ABORTED_CONNECTS":
		abortedConnections, err = strconv.Atoi(value)

		// If we don't have a previous value for the total aborted connections
		// then acps is technically 0 as we don't know it yet
		if previous == nil || previous.Metrics.AbortedConnections == 0 {
			status.Metrics.AbortedConnectionsPerSecond = 0
			status.Metrics.AbortedConnections = abortedConnections
		// Otherwise the value of acps is the diff between the current
		// and previous count of connections
		} else {
			diff = abortedConnections - previous.Metrics.AbortedConnections

			// qps can never be below 0..
			if diff > 0 {
				status.Metrics.AbortedConnectionsPerSecond = diff
			} else {
				status.Metrics.AbortedConnectionsPerSecond = 0
			}

			status.Metrics.AbortedConnections = abortedConnections
		}
	// Queries per second
	case "QUERIES":
		queries, err = strconv.Atoi(value)

		// If we don't have a previous value for the total queries
		// then qps is technically 0 as we don't know it yet
		if previous == nil || previous.Metrics.Queries == 0 {
			status.Metrics.QueriesPerSecond = 0
			status.Metrics.Queries = queries
		// Otherwise the value of qps is the diff between the current
		// and previous count of queries
		} else {
			diff = queries - previous.Metrics.Queries

			// qps can never be below 0..
			if diff > 0 {
				status.Metrics.QueriesPerSecond = diff
			} else {
				status.Metrics.QueriesPerSecond = 0
			}

			status.Metrics.Queries = queries
		}
	// Reads per second
	case "COM_SELECT", "COM_INSERT_SELECT", "COM_REPLACE_SELECT":
		reads, err = strconv.Atoi(value)

		// If we don't have a previous value for the total reads
		// then rps is technically 0 as we don't know it yet
		if previous == nil || previous.Metrics.Reads == 0 {
			status.Metrics.ReadsPerSecond = 0
			status.Metrics.Reads = reads
		// Otherwise the value of rps is the diff between the current
		// and previous count of reads
		} else {
			diff = reads - previous.Metrics.Reads

			// rps can never be below 0..
			if diff > 0 {
				status.Metrics.ReadsPerSecond = diff
			} else {
				status.Metrics.ReadsPerSecond = 0
			}

			status.Metrics.Reads = reads
		}
	// Uptime
	case "UPTIME":
		uptime, err = strconv.Atoi(value)
		status.Metrics.Uptime = uptime
	}

	if err != nil {
		return status, err
	} else {
		return status, nil
	}
}

// Process variables returned from the GLOBAL_VARIABLES table
func processVariables(status *DatabaseStatus, key string, value string) (*DatabaseStatus, error) {
	var (
		err            error
		maxConnections int
	)

	// Max allowed connections
	if key == "MAX_CONNECTIONS" {
		maxConnections, err = strconv.Atoi(value)
		status.Variables.MaxConnections = maxConnections
	}

	if err != nil {
		return status, err
	} else {
		return status, nil
	}
}
<|endoftext|>


<commit_before>package conveyor

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strings"
	"sync"
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/fsouza/go-dockerclient"
	"github.com/google/go-github/github"
	"github.com/remind101/pkg/reporter"

	"golang.org/x/net/context"
)

const (
	// Context is used for the commit status context.
	Context = "container/docker"

	// DefaultBuilderImage is the docker image used to build docker images.
	DefaultBuilderImage = "remind101/conveyor-builder"

	// DefaultDataVolume is the default name of a container serving as a
	// data volume for ssh keys and docker credentials. In general, you
	// shouldn't need to change this.
	DefaultDataVolume = "data"

	// DefaultTimeout is the default amount of time to wait for a build
	// to complete before cancelling it.
	DefaultTimeout = 20 * time.Minute
)

var (
	// ErrShuttingDown can be returned by builders if they're shutting down
	// and not accepting more jobs.
	ErrShuttingDown = errors.New("shutting down")
)

// BuildCanceledError is returned if the build is canceled, or times out and the
// container returns an error.
type BuildCanceledError struct {
	Err error
}

// Error implements the error interface.
func (e *BuildCanceledError) Error() string {
	return e.Err.Error()
}

type BuildOptions struct {
	// Repository is the repo to build.
	Repository string
	// Sha is the git commit to build.
	Sha string
	// Branch is the name of the branch that this build relates to.
	Branch string
	// Set to true to disable the layer cache. The zero value is to enable
	// caching.
	NoCache bool
}

// Builder represents something that can build a Docker image.
type Builder interface {
	// Build should build the docker image, tag it and push it to the docker
	// registry. This should return the sha256 digest of the image.
	Build(context.Context, Logger, BuildOptions) (string, error)
}

// BuilderFunc is a function that implements the Builder interface.
type BuilderFunc func(context.Context, Logger, BuildOptions) (string, error)

func (fn BuilderFunc) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	return fn(ctx, w, opts)
}

// Conveyor serves as a builder.
type Conveyor struct {
	Builder

	// A Reporter to use to report errors.
	Reporter reporter.Reporter

	// Timeout controls how long to wait before canceling a build. A timeout
	// of 0 means no timeout.
	Timeout time.Duration
}

// New returns a new Conveyor instance.
func New(b Builder) *Conveyor {
	return &Conveyor{
		Builder: WithCancel(b),
		Timeout: DefaultTimeout,
	}
}

func (c *Conveyor) Build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	log.Printf("Starting build: repository=%s branch=%s sha=%s",
		opts.Repository,
		opts.Branch,
		opts.Sha,
	)

	// Embed the reporter in the context.Context.
	ctx = reporter.WithReporter(ctx, c.reporter())

	if c.Timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, c.Timeout)
		defer cancel() // Release resources.
	}

	reporter.AddContext(ctx, "options", opts)
	defer reporter.Monitor(ctx)

	defer func() {
		if err != nil {
			reporter.Report(ctx, err)
		}
	}()

	id, err = c.build(ctx, w, opts)
	return
}

// Build performs the build and ensures that the output stream is closed.
func (c *Conveyor) build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	defer func() {
		var closeErr error
		if w != nil {
			closeErr = w.Close()
		}
		if err == nil {
			// If there was no error from the builder, let the
			// downstream know that there was an error closing the
			// output stream.
			err = closeErr
		}
	}()

	id, err = c.Builder.Build(ctx, w, opts)
	return
}

func (c *Conveyor) Cancel() error {
	if b, ok := c.Builder.(*CancelBuilder); ok {
		return b.Cancel()
	}

	return fmt.Errorf("Builder does not support Cancel()")
}

func (c *Conveyor) reporter() reporter.Reporter {
	if c.Reporter == nil {
		return reporter.ReporterFunc(func(ctx context.Context, err error) error {
			fmt.Fprintf(os.Stderr, "reporting err: %v", err)
			return nil
		})
	}

	return c.Reporter
}

// DockerBuilder is a Builder implementation that runs the build in a docker
// container.
type DockerBuilder struct {
	// dataVolume is the name of the volume that contains ssh keys and
	// configuration data.
	DataVolume string
	// Name of the image to use to build the docker image. Defaults to
	// DefaultBuilderImage.
	Image string
	// Set to true to enable dry runs. This sets the `DRY` environment
	// variable within the builder container to `true`. The behavior of this
	// flag depends on how the builder image handles the `DRY` environment
	// variable.
	DryRun bool

	client *docker.Client
}

// NewDockerBuilder returns a new DockerBuilder backed by the docker client.
func NewDockerBuilder(c *docker.Client) *DockerBuilder {
	return &DockerBuilder{client: c}
}

// NewDockerBuilderFromEnv returns a new DockerBuilder with a docker client
// configured from the standard Docker environment variables.
func NewDockerBuilderFromEnv() (*DockerBuilder, error) {
	c, err := docker.NewClientFromEnv()
	if err != nil {
		return nil, err
	}

	return NewDockerBuilder(c), nil
}

// Build executes the docker image.
func (b *DockerBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	env := []string{
		fmt.Sprintf("REPOSITORY=%s", opts.Repository),
		fmt.Sprintf("BRANCH=%s", opts.Branch),
		fmt.Sprintf("SHA=%s", opts.Sha),
		fmt.Sprintf("DRY=%s", b.dryRun()),
		fmt.Sprintf("CACHE=%s", b.cache(opts)),
	}

	name := strings.Join([]string{
		strings.Replace(opts.Repository, "/", "-", -1),
		opts.Sha,
		uuid.New(),
	}, "-")

	c, err := b.client.CreateContainer(docker.CreateContainerOptions{
		Name: name,
		Config: &docker.Config{
			Tty:          true,
			AttachStdout: true,
			AttachStderr: true,
			OpenStdin:    true,
			Image:        b.image(),
			Hostname:     hostname,
			Env:          env,
		},
	})
	if err != nil {
		return "", fmt.Errorf("create container: %v", err)
	}
	defer b.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            c.ID,
		RemoveVolumes: true,
		Force:         true,
	})

	reporter.AddContext(ctx, "container_id", c.ID)

	if err := b.client.StartContainer(c.ID, &docker.HostConfig{
		Privileged:  true,
		VolumesFrom: []string{b.dataVolume()},
	}); err != nil {
		return "", fmt.Errorf("start container: %v", err)
	}

	done := make(chan error, 1)
	go func() {
		err := b.client.AttachToContainer(docker.AttachToContainerOptions{
			Container:    c.ID,
			OutputStream: w,
			ErrorStream:  w,
			Logs:         true,
			Stream:       true,
			Stdout:       true,
			Stderr:       true,
			RawTerminal:  true,
		})
		done <- err
	}()

	var canceled bool
	select {
	case <-ctx.Done():
		// Build was canceled or the build timedout. Stop the container
		// prematurely. We'll SIGTERM and give it 10 seconds to stop,
		// after that we'll SIGKILL.
		if err := b.client.StopContainer(c.ID, 10); err != nil {
			return "", fmt.Errorf("stop: %v", err)
		}

		// Wait for log streaming to finish.
		err := <-done
		if err != nil {
			return "", fmt.Errorf("attach: %v", err)
		}

		canceled = true
	case err := <-done:
		if err != nil {
			return "", fmt.Errorf("attach: %v", err)
		}
	}

	exit, err := b.client.WaitContainer(c.ID)
	if err != nil {
		return "", fmt.Errorf("wait container: %v", err)
	}

	// A non-zero exit status means the build failed.
	if exit != 0 {
		err := fmt.Errorf("container returned a non-zero exit code: %d", exit)
		if canceled {
			err = &BuildCanceledError{
				Err: err,
			}
		}
		return "", err
	}

	// TODO: Return sha256
	return "", nil
}

func (b *DockerBuilder) dryRun() string {
	if b.DryRun {
		return "true"
	}
	return ""
}

func (b *DockerBuilder) image() string {
	if b.Image == "" {
		return DefaultBuilderImage
	}
	return b.Image
}

func (b *DockerBuilder) dataVolume() string {
	if b.DataVolume == "" {
		return DefaultDataVolume
	}
	return b.DataVolume
}

func (b *DockerBuilder) cache(opts BuildOptions) string {
	if opts.NoCache {
		return "off"
	}

	return "on"
}

// statusUpdaterBuilder is a Builder implementation that updates the commit
// status in github.
type statusUpdaterBuilder struct {
	Builder
	github GitHubClient
	since  func(time.Time) time.Duration
}

// UpdateGitHubCommitStatus wraps b to update the GitHub commit status when a
// build starts, and stops.
func UpdateGitHubCommitStatus(b Builder, g GitHubClient) *statusUpdaterBuilder {
	return &statusUpdaterBuilder{
		Builder: b,
		github:  g,
		since:   time.Since,
	}
}

func (b *statusUpdaterBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	t := time.Now()

	defer func() {
		duration := b.since(t)
		description := fmt.Sprintf("Image built in %v.", duration)
		status := "success"
		if err != nil {
			status = "failure"
			description = err.Error()
		}
		b.updateStatus(w, opts, status, description)
	}()

	if err = b.updateStatus(w, opts, "pending", "Image building."); err != nil {
		return
	}

	id, err = b.Builder.Build(ctx, w, opts)
	return
}

// updateStatus updates the given commit with a new status.
func (b *statusUpdaterBuilder) updateStatus(w Logger, opts BuildOptions, status string, description string) error {
	context := Context
	parts := strings.SplitN(opts.Repository, "/", 2)

	var desc *string
	if description != "" {
		desc = &description
	}
	var url *string
	if status == "success" || status == "failure" || status == "error" {
		url = github.String(w.URL())
	}

	_, _, err := b.github.CreateStatus(parts[0], parts[1], opts.Sha, &github.RepoStatus{
		State:       &status,
		Context:     &context,
		Description: desc,
		TargetURL:   url,
	})
	return err
}

// BuildAsync wraps a Builder to run the build in a goroutine.
func BuildAsync(b Builder) Builder {
	build := func(ctx context.Context, w Logger, opts BuildOptions) {
		if _, err := b.Build(ctx, w, opts); err != nil {
			log.Printf("build err: %v", err)
		}
	}

	return BuilderFunc(func(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
		go build(ctx, w, opts)
		return "", nil
	})
}

// WithCancel wraps a Builder with a method to stop all builds.
func WithCancel(b Builder) *CancelBuilder {
	return &CancelBuilder{
		Builder:  b,
		shutdown: make(chan struct{}),
		builds:   make(map[context.Context]context.CancelFunc),
	}
}

type CancelBuilder struct {
	Builder

	shutdown chan struct{}

	sync.Mutex
	stopped bool
	builds  map[context.Context]context.CancelFunc
}

func (b *CancelBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	if b.stopped {
		return "", ErrShuttingDown
	}

	ctx = b.addBuild(ctx)
	defer b.removeBuild(ctx)

	return b.Builder.Build(ctx, w, opts)
}

func (b *CancelBuilder) Cancel() error {
	b.Lock()
	defer b.Unlock()

	// Mark as stopped so we don't accept anymore builds.
	b.stopped = true

	// Cancel each build.
	for _, cancel := range b.builds {
		cancel()
	}

	// Wait for each build to complete.
	for ctx := range b.builds {
		<-ctx.Done()
	}

	return nil
}

func (b *CancelBuilder) addBuild(ctx context.Context) context.Context {
	b.Lock()
	defer b.Unlock()

	ctx, cancel := context.WithCancel(ctx)
	b.builds[ctx] = cancel
	return ctx
}

func (b *CancelBuilder) removeBuild(ctx context.Context) {
	b.Lock()
	defer b.Unlock()

	delete(b.builds, ctx)
}

var hostname string

func init() {
	hostname, _ = os.Hostname()
}
<commit_msg>Append (canceled) to BuildCanceledError's.<commit_after>package conveyor

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strings"
	"sync"
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/fsouza/go-dockerclient"
	"github.com/google/go-github/github"
	"github.com/remind101/pkg/reporter"

	"golang.org/x/net/context"
)

const (
	// Context is used for the commit status context.
	Context = "container/docker"

	// DefaultBuilderImage is the docker image used to build docker images.
	DefaultBuilderImage = "remind101/conveyor-builder"

	// DefaultDataVolume is the default name of a container serving as a
	// data volume for ssh keys and docker credentials. In general, you
	// shouldn't need to change this.
	DefaultDataVolume = "data"

	// DefaultTimeout is the default amount of time to wait for a build
	// to complete before cancelling it.
	DefaultTimeout = 20 * time.Minute
)

var (
	// ErrShuttingDown can be returned by builders if they're shutting down
	// and not accepting more jobs.
	ErrShuttingDown = errors.New("shutting down")
)

// BuildCanceledError is returned if the build is canceled, or times out and the
// container returns an error.
type BuildCanceledError struct {
	Err error
}

// Error implements the error interface.
func (e *BuildCanceledError) Error() string {
	return fmt.Sprintf("%s (canceled)", e.Err.Error())
}

type BuildOptions struct {
	// Repository is the repo to build.
	Repository string
	// Sha is the git commit to build.
	Sha string
	// Branch is the name of the branch that this build relates to.
	Branch string
	// Set to true to disable the layer cache. The zero value is to enable
	// caching.
	NoCache bool
}

// Builder represents something that can build a Docker image.
type Builder interface {
	// Build should build the docker image, tag it and push it to the docker
	// registry. This should return the sha256 digest of the image.
	Build(context.Context, Logger, BuildOptions) (string, error)
}

// BuilderFunc is a function that implements the Builder interface.
type BuilderFunc func(context.Context, Logger, BuildOptions) (string, error)

func (fn BuilderFunc) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	return fn(ctx, w, opts)
}

// Conveyor serves as a builder.
type Conveyor struct {
	Builder

	// A Reporter to use to report errors.
	Reporter reporter.Reporter

	// Timeout controls how long to wait before canceling a build. A timeout
	// of 0 means no timeout.
	Timeout time.Duration
}

// New returns a new Conveyor instance.
func New(b Builder) *Conveyor {
	return &Conveyor{
		Builder: WithCancel(b),
		Timeout: DefaultTimeout,
	}
}

func (c *Conveyor) Build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	log.Printf("Starting build: repository=%s branch=%s sha=%s",
		opts.Repository,
		opts.Branch,
		opts.Sha,
	)

	// Embed the reporter in the context.Context.
	ctx = reporter.WithReporter(ctx, c.reporter())

	if c.Timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, c.Timeout)
		defer cancel() // Release resources.
	}

	reporter.AddContext(ctx, "options", opts)
	defer reporter.Monitor(ctx)

	defer func() {
		if err != nil {
			reporter.Report(ctx, err)
		}
	}()

	id, err = c.build(ctx, w, opts)
	return
}

// Build performs the build and ensures that the output stream is closed.
func (c *Conveyor) build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	defer func() {
		var closeErr error
		if w != nil {
			closeErr = w.Close()
		}
		if err == nil {
			// If there was no error from the builder, let the
			// downstream know that there was an error closing the
			// output stream.
			err = closeErr
		}
	}()

	id, err = c.Builder.Build(ctx, w, opts)
	return
}

func (c *Conveyor) Cancel() error {
	if b, ok := c.Builder.(*CancelBuilder); ok {
		return b.Cancel()
	}

	return fmt.Errorf("Builder does not support Cancel()")
}

func (c *Conveyor) reporter() reporter.Reporter {
	if c.Reporter == nil {
		return reporter.ReporterFunc(func(ctx context.Context, err error) error {
			fmt.Fprintf(os.Stderr, "reporting err: %v", err)
			return nil
		})
	}

	return c.Reporter
}

// DockerBuilder is a Builder implementation that runs the build in a docker
// container.
type DockerBuilder struct {
	// dataVolume is the name of the volume that contains ssh keys and
	// configuration data.
	DataVolume string
	// Name of the image to use to build the docker image. Defaults to
	// DefaultBuilderImage.
	Image string
	// Set to true to enable dry runs. This sets the `DRY` environment
	// variable within the builder container to `true`. The behavior of this
	// flag depends on how the builder image handles the `DRY` environment
	// variable.
	DryRun bool

	client *docker.Client
}

// NewDockerBuilder returns a new DockerBuilder backed by the docker client.
func NewDockerBuilder(c *docker.Client) *DockerBuilder {
	return &DockerBuilder{client: c}
}

// NewDockerBuilderFromEnv returns a new DockerBuilder with a docker client
// configured from the standard Docker environment variables.
func NewDockerBuilderFromEnv() (*DockerBuilder, error) {
	c, err := docker.NewClientFromEnv()
	if err != nil {
		return nil, err
	}

	return NewDockerBuilder(c), nil
}

// Build executes the docker image.
func (b *DockerBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	env := []string{
		fmt.Sprintf("REPOSITORY=%s", opts.Repository),
		fmt.Sprintf("BRANCH=%s", opts.Branch),
		fmt.Sprintf("SHA=%s", opts.Sha),
		fmt.Sprintf("DRY=%s", b.dryRun()),
		fmt.Sprintf("CACHE=%s", b.cache(opts)),
	}

	name := strings.Join([]string{
		strings.Replace(opts.Repository, "/", "-", -1),
		opts.Sha,
		uuid.New(),
	}, "-")

	c, err := b.client.CreateContainer(docker.CreateContainerOptions{
		Name: name,
		Config: &docker.Config{
			Tty:          true,
			AttachStdout: true,
			AttachStderr: true,
			OpenStdin:    true,
			Image:        b.image(),
			Hostname:     hostname,
			Env:          env,
		},
	})
	if err != nil {
		return "", fmt.Errorf("create container: %v", err)
	}
	defer b.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            c.ID,
		RemoveVolumes: true,
		Force:         true,
	})

	reporter.AddContext(ctx, "container_id", c.ID)

	if err := b.client.StartContainer(c.ID, &docker.HostConfig{
		Privileged:  true,
		VolumesFrom: []string{b.dataVolume()},
	}); err != nil {
		return "", fmt.Errorf("start container: %v", err)
	}

	done := make(chan error, 1)
	go func() {
		err := b.client.AttachToContainer(docker.AttachToContainerOptions{
			Container:    c.ID,
			OutputStream: w,
			ErrorStream:  w,
			Logs:         true,
			Stream:       true,
			Stdout:       true,
			Stderr:       true,
			RawTerminal:  true,
		})
		done <- err
	}()

	var canceled bool
	select {
	case <-ctx.Done():
		// Build was canceled or the build timedout. Stop the container
		// prematurely. We'll SIGTERM and give it 10 seconds to stop,
		// after that we'll SIGKILL.
		if err := b.client.StopContainer(c.ID, 10); err != nil {
			return "", fmt.Errorf("stop: %v", err)
		}

		// Wait for log streaming to finish.
		err := <-done
		if err != nil {
			return "", fmt.Errorf("attach: %v", err)
		}

		canceled = true
	case err := <-done:
		if err != nil {
			return "", fmt.Errorf("attach: %v", err)
		}
	}

	exit, err := b.client.WaitContainer(c.ID)
	if err != nil {
		return "", fmt.Errorf("wait container: %v", err)
	}

	// A non-zero exit status means the build failed.
	if exit != 0 {
		err := fmt.Errorf("container returned a non-zero exit code: %d", exit)
		if canceled {
			err = &BuildCanceledError{
				Err: err,
			}
		}
		return "", err
	}

	// TODO: Return sha256
	return "", nil
}

func (b *DockerBuilder) dryRun() string {
	if b.DryRun {
		return "true"
	}
	return ""
}

func (b *DockerBuilder) image() string {
	if b.Image == "" {
		return DefaultBuilderImage
	}
	return b.Image
}

func (b *DockerBuilder) dataVolume() string {
	if b.DataVolume == "" {
		return DefaultDataVolume
	}
	return b.DataVolume
}

func (b *DockerBuilder) cache(opts BuildOptions) string {
	if opts.NoCache {
		return "off"
	}

	return "on"
}

// statusUpdaterBuilder is a Builder implementation that updates the commit
// status in github.
type statusUpdaterBuilder struct {
	Builder
	github GitHubClient
	since  func(time.Time) time.Duration
}

// UpdateGitHubCommitStatus wraps b to update the GitHub commit status when a
// build starts, and stops.
func UpdateGitHubCommitStatus(b Builder, g GitHubClient) *statusUpdaterBuilder {
	return &statusUpdaterBuilder{
		Builder: b,
		github:  g,
		since:   time.Since,
	}
}

func (b *statusUpdaterBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	t := time.Now()

	defer func() {
		duration := b.since(t)
		description := fmt.Sprintf("Image built in %v.", duration)
		status := "success"
		if err != nil {
			status = "failure"
			description = err.Error()
		}
		b.updateStatus(w, opts, status, description)
	}()

	if err = b.updateStatus(w, opts, "pending", "Image building."); err != nil {
		return
	}

	id, err = b.Builder.Build(ctx, w, opts)
	return
}

// updateStatus updates the given commit with a new status.
func (b *statusUpdaterBuilder) updateStatus(w Logger, opts BuildOptions, status string, description string) error {
	context := Context
	parts := strings.SplitN(opts.Repository, "/", 2)

	var desc *string
	if description != "" {
		desc = &description
	}
	var url *string
	if status == "success" || status == "failure" || status == "error" {
		url = github.String(w.URL())
	}

	_, _, err := b.github.CreateStatus(parts[0], parts[1], opts.Sha, &github.RepoStatus{
		State:       &status,
		Context:     &context,
		Description: desc,
		TargetURL:   url,
	})
	return err
}

// BuildAsync wraps a Builder to run the build in a goroutine.
func BuildAsync(b Builder) Builder {
	build := func(ctx context.Context, w Logger, opts BuildOptions) {
		if _, err := b.Build(ctx, w, opts); err != nil {
			log.Printf("build err: %v", err)
		}
	}

	return BuilderFunc(func(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
		go build(ctx, w, opts)
		return "", nil
	})
}

// WithCancel wraps a Builder with a method to stop all builds.
func WithCancel(b Builder) *CancelBuilder {
	return &CancelBuilder{
		Builder:  b,
		shutdown: make(chan struct{}),
		builds:   make(map[context.Context]context.CancelFunc),
	}
}

type CancelBuilder struct {
	Builder

	shutdown chan struct{}

	sync.Mutex
	stopped bool
	builds  map[context.Context]context.CancelFunc
}

func (b *CancelBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	if b.stopped {
		return "", ErrShuttingDown
	}

	ctx = b.addBuild(ctx)
	defer b.removeBuild(ctx)

	return b.Builder.Build(ctx, w, opts)
}

func (b *CancelBuilder) Cancel() error {
	b.Lock()
	defer b.Unlock()

	// Mark as stopped so we don't accept anymore builds.
	b.stopped = true

	// Cancel each build.
	for _, cancel := range b.builds {
		cancel()
	}

	// Wait for each build to complete.
	for ctx := range b.builds {
		<-ctx.Done()
	}

	return nil
}

func (b *CancelBuilder) addBuild(ctx context.Context) context.Context {
	b.Lock()
	defer b.Unlock()

	ctx, cancel := context.WithCancel(ctx)
	b.builds[ctx] = cancel
	return ctx
}

func (b *CancelBuilder) removeBuild(ctx context.Context) {
	b.Lock()
	defer b.Unlock()

	delete(b.builds, ctx)
}

var hostname string

func init() {
	hostname, _ = os.Hostname()
}
<|endoftext|>