<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/alecthomas/kingpin"
	"github.com/estafette/estafette-ci-contracts"
	crypt "github.com/estafette/estafette-ci-crypt"
	manifest "github.com/estafette/estafette-ci-manifest"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"

	stdlog "log"
)

var (
	version   string
	branch    string
	revision  string
	buildDate string
	goVersion = runtime.Version()

	secretDecryptionKey = kingpin.Flag("secret-decryption-key", "The AES-256 key used to decrypt secrets that have been encrypted with it.").Envar("SECRET_DECRYPTION_KEY").String()
	runAsJob            = kingpin.Flag("run-as-job", "To run the builder as a job and prevent build failures to fail the job.").Default("true").OverrideDefaultFromEnvar("RUN_AS_JOB").Bool()
)

func main() {

	// parse command line parameters
	kingpin.Parse()

	// bootstrap
	secretHelper := crypt.NewSecretHelper(*secretDecryptionKey)
	envvarHelper := NewEnvvarHelper("ESTAFETTE_", secretHelper)
	whenEvaluator := NewWhenEvaluator(envvarHelper)
	obfuscator := NewObfuscator(secretHelper)
	dockerRunner := NewDockerRunner(envvarHelper, obfuscator)
	pipelineRunner := NewPipelineRunner(envvarHelper, whenEvaluator, dockerRunner)
	endOfLifeHelper := NewEndOfLifeHelper(envvarHelper, *runAsJob)

	// detect controlling server
	ciServer := envvarHelper.getEstafetteEnv("ESTAFETTE_CI_SERVER")

	if ciServer == "gocd" {

		// pretty print for go.cd integration
		log.Logger = zerolog.New(zerolog.ConsoleWriter{Out: os.Stderr}).With().
			Timestamp().
			Logger()

		stdlog.SetFlags(0)
		stdlog.SetOutput(log.Logger)

		// log startup message
		log.Info().
			Str("branch", branch).
			Str("revision", revision).
			Str("buildDate", buildDate).
			Str("goVersion", goVersion).
			Msg("Starting estafette-ci-builder...")

		// create docker client
		_, err := dockerRunner.createDockerClient()
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Failed creating a docker client")
		}

		// read yaml
		manifest, err := manifest.ReadManifestFromFile(".estafette.yaml")
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Reading .estafette.yaml manifest failed")
		}

		// initialize obfuscator
		err = obfuscator.CollectSecrets(manifest)
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Collecting secrets to obfuscate failed")
		}

		// get current working directory
		dir, err := os.Getwd()
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Getting current working directory failed")
		}

		log.Info().Msgf("Running %v stages", len(manifest.Stages))

		err = envvarHelper.setEstafetteGlobalEnvvars()
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Setting global environment variables failed")
		}

		// collect estafette and 'global' envvars from manifest
		estafetteEnvvars := envvarHelper.collectEstafetteEnvvars(manifest)
		globalEnvvars := envvarHelper.collectGlobalEnvvars(manifest)

		// merge estafette and global envvars
		envvars := envvarHelper.overrideEnvvars(estafetteEnvvars, globalEnvvars)

		result, err := pipelineRunner.runStages(manifest.Stages, dir, envvars)
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Executing stages from manifest failed")
		}

		renderStats(result)

		handleExit(result)

	} else if ciServer == "estafette" {

		// log as severity for stackdriver logging to recognize the level
		zerolog.LevelFieldName = "severity"

		gitName := envvarHelper.getEstafetteEnv("ESTAFETTE_GIT_NAME")
		gitBranch := envvarHelper.getEstafetteEnv("ESTAFETTE_GIT_BRANCH")
		gitRevision := envvarHelper.getEstafetteEnv("ESTAFETTE_GIT_REVISION")
		jobName := envvarHelper.getEstafetteEnv("ESTAFETTE_BUILD_JOB_NAME")
		builderTrack := envvarHelper.getEstafetteEnv("ESTAFETTE_CI_BUILDER_TRACK")
		if builderTrack == "" {
			builderTrack = "stable"
		}
		buildVersion := envvarHelper.getEstafetteEnv("ESTAFETTE_BUILD_VERSION")
		releaseName := envvarHelper.getEstafetteEnv("ESTAFETTE_RELEASE_NAME")
		releaseIDValue := envvarHelper.getEstafetteEnv("ESTAFETTE_RELEASE_ID")
		releaseID, _ := strconv.Atoi(releaseIDValue)

		buildLog := contracts.BuildLog{
			RepoSource:   envvarHelper.getEstafetteEnv("ESTAFETTE_GIT_SOURCE"),
			RepoOwner:    strings.Split(gitName, "/")[0],
			RepoName:     strings.Split(gitName, "/")[1],
			RepoBranch:   gitBranch,
			RepoRevision: gitRevision,
			Steps:        make([]contracts.BuildLogStep, 0),
		}

		// log to file and stdout
		logFile, err := os.OpenFile("/log.txt", os.O_CREATE|os.O_APPEND|os.O_RDWR, 0666)
		if err != nil {
			log.Fatal().Err(err).Msg("Failed to create log file log.txt")
		}
		defer logFile.Close()
		multiLogWriter := io.MultiWriter(os.Stdout, logFile)

		// set some default fields added to all logs
		log.Logger = zerolog.New(multiLogWriter).With().
			Timestamp().
			Str("app", "estafette-ci-builder").
			Str("version", version).
			Str("jobName", jobName).
			Str("gitName", gitName).
			Str("gitBranch", gitBranch).
			Str("gitRevision", gitRevision).
			Logger()

		stdlog.SetFlags(0)
		stdlog.SetOutput(log.Logger)

		// log startup message
		log.Info().
			Str("branch", branch).
			Str("revision", revision).
			Str("buildDate", buildDate).
			Str("goVersion", goVersion).
			Msg("Starting estafette-ci-builder...")

		// start docker daemon
		err = dockerRunner.startDockerDaemon()
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Error starting docker daemon")
		}

		// wait for docker daemon to be ready for usage
		dockerRunner.waitForDockerDaemon()

		// get current working directory
		dir, err := os.Getwd()
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Getting current working directory failed")
		}

		// set some envvars
		err = envvarHelper.setEstafetteGlobalEnvvars()
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Setting global environment variables failed")
		}

		// get manifest from envvar and unmarshal
		manifestJSON := os.Getenv("ESTAFETTE_CI_MANIFEST_JSON")
		var manifest manifest.EstafetteManifest
		json.Unmarshal([]byte(manifestJSON), &manifest)

		// initialize obfuscator
		err = obfuscator.CollectSecrets(manifest)
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Collecting secrets to obfuscate failed")
		}

		// check whether this is a regular build or a release
		stages := manifest.Stages
		if releaseID > 0 {
			// check if the release is defined
			releaseExists := false
			for _, r := range manifest.Releases {
				if r.Name == releaseName {
					releaseExists = true
					stages = r.Stages
				}
			}
			if !releaseExists {
				endOfLifeHelper.handleFatal(buildLog, err, fmt.Sprintf("Release %v does not exist", releaseName))
			}
			log.Info().Msgf("Starting release %v at version %v...", releaseName, buildVersion)
		} else {
			log.Info().Msgf("Starting build version %v...", buildVersion)
		}

		// create docker client
		_, err = dockerRunner.createDockerClient()
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Failed creating a docker client")
		}

		// get private container registries credentials
		credentialsJSON := os.Getenv("ESTAFETTE_CI_REPOSITORY_CREDENTIALS_JSON")
		if credentialsJSON != "" {
			var credentials []*contracts.ContainerRepositoryCredentialConfig
			json.Unmarshal([]byte(credentialsJSON), &credentials)
			dockerRunner.setRepositoryCredentials(credentials)
		}

		// collect estafette envvars and run stages from manifest
		log.Info().Msgf("Running %v stages", len(stages))
		estafetteEnvvars := envvarHelper.collectEstafetteEnvvars(manifest)
		globalEnvvars := envvarHelper.collectGlobalEnvvars(manifest)
		envvars := envvarHelper.overrideEnvvars(estafetteEnvvars, globalEnvvars)
		result, err := pipelineRunner.runStages(stages, dir, envvars)
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Executing stages from manifest failed")
		}

		// send result to ci-api
		log.Info().Interface("result", result).Msg("Finished running stages")
		buildLog.Steps = transformPipelineRunResultToBuildLogSteps(result)
		endOfLifeHelper.sendBuildJobLogEvent(buildLog)
		buildStatus := "succeeded"
		if result.HasErrors() {
			buildStatus = "failed"
		}
		endOfLifeHelper.sendBuildFinishedEvent(buildStatus)
		os.Exit(0)
	} else {
		// Set up a simple console logger
		log.Logger = zerolog.New(zerolog.ConsoleWriter{Out: os.Stderr}).With().
			Timestamp().
			Logger()

		log.Warn().Msgf("The CI Server (\"%s\") is not recognized, exiting.", ciServer)
	}
}
<commit_msg>don't swallow failure if RUN_AS_JOB=false<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/alecthomas/kingpin"
	"github.com/estafette/estafette-ci-contracts"
	crypt "github.com/estafette/estafette-ci-crypt"
	manifest "github.com/estafette/estafette-ci-manifest"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"

	stdlog "log"
)

var (
	version   string
	branch    string
	revision  string
	buildDate string
	goVersion = runtime.Version()

	secretDecryptionKey = kingpin.Flag("secret-decryption-key", "The AES-256 key used to decrypt secrets that have been encrypted with it.").Envar("SECRET_DECRYPTION_KEY").String()
	runAsJob            = kingpin.Flag("run-as-job", "To run the builder as a job and prevent build failures to fail the job.").Default("true").OverrideDefaultFromEnvar("RUN_AS_JOB").Bool()
)

func main() {

	// parse command line parameters
	kingpin.Parse()

	// bootstrap
	secretHelper := crypt.NewSecretHelper(*secretDecryptionKey)
	envvarHelper := NewEnvvarHelper("ESTAFETTE_", secretHelper)
	whenEvaluator := NewWhenEvaluator(envvarHelper)
	obfuscator := NewObfuscator(secretHelper)
	dockerRunner := NewDockerRunner(envvarHelper, obfuscator)
	pipelineRunner := NewPipelineRunner(envvarHelper, whenEvaluator, dockerRunner)
	endOfLifeHelper := NewEndOfLifeHelper(envvarHelper, *runAsJob)

	// detect controlling server
	ciServer := envvarHelper.getEstafetteEnv("ESTAFETTE_CI_SERVER")

	if ciServer == "gocd" {

		// pretty print for go.cd integration
		log.Logger = zerolog.New(zerolog.ConsoleWriter{Out: os.Stderr}).With().
			Timestamp().
			Logger()

		stdlog.SetFlags(0)
		stdlog.SetOutput(log.Logger)

		// log startup message
		log.Info().
			Str("branch", branch).
			Str("revision", revision).
			Str("buildDate", buildDate).
			Str("goVersion", goVersion).
			Msg("Starting estafette-ci-builder...")

		// create docker client
		_, err := dockerRunner.createDockerClient()
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Failed creating a docker client")
		}

		// read yaml
		manifest, err := manifest.ReadManifestFromFile(".estafette.yaml")
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Reading .estafette.yaml manifest failed")
		}

		// initialize obfuscator
		err = obfuscator.CollectSecrets(manifest)
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Collecting secrets to obfuscate failed")
		}

		// get current working directory
		dir, err := os.Getwd()
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Getting current working directory failed")
		}

		log.Info().Msgf("Running %v stages", len(manifest.Stages))

		err = envvarHelper.setEstafetteGlobalEnvvars()
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Setting global environment variables failed")
		}

		// collect estafette and 'global' envvars from manifest
		estafetteEnvvars := envvarHelper.collectEstafetteEnvvars(manifest)
		globalEnvvars := envvarHelper.collectGlobalEnvvars(manifest)

		// merge estafette and global envvars
		envvars := envvarHelper.overrideEnvvars(estafetteEnvvars, globalEnvvars)

		result, err := pipelineRunner.runStages(manifest.Stages, dir, envvars)
		if err != nil {
			endOfLifeHelper.handleGocdFatal(err, "Executing stages from manifest failed")
		}

		renderStats(result)

		handleExit(result)

	} else if ciServer == "estafette" {

		// log as severity for stackdriver logging to recognize the level
		zerolog.LevelFieldName = "severity"

		gitName := envvarHelper.getEstafetteEnv("ESTAFETTE_GIT_NAME")
		gitBranch := envvarHelper.getEstafetteEnv("ESTAFETTE_GIT_BRANCH")
		gitRevision := envvarHelper.getEstafetteEnv("ESTAFETTE_GIT_REVISION")
		jobName := envvarHelper.getEstafetteEnv("ESTAFETTE_BUILD_JOB_NAME")
		builderTrack := envvarHelper.getEstafetteEnv("ESTAFETTE_CI_BUILDER_TRACK")
		if builderTrack == "" {
			builderTrack = "stable"
		}
		buildVersion := envvarHelper.getEstafetteEnv("ESTAFETTE_BUILD_VERSION")
		releaseName := envvarHelper.getEstafetteEnv("ESTAFETTE_RELEASE_NAME")
		releaseIDValue := envvarHelper.getEstafetteEnv("ESTAFETTE_RELEASE_ID")
		releaseID, _ := strconv.Atoi(releaseIDValue)

		buildLog := contracts.BuildLog{
			RepoSource:   envvarHelper.getEstafetteEnv("ESTAFETTE_GIT_SOURCE"),
			RepoOwner:    strings.Split(gitName, "/")[0],
			RepoName:     strings.Split(gitName, "/")[1],
			RepoBranch:   gitBranch,
			RepoRevision: gitRevision,
			Steps:        make([]contracts.BuildLogStep, 0),
		}

		// log to file and stdout
		logFile, err := os.OpenFile("/log.txt", os.O_CREATE|os.O_APPEND|os.O_RDWR, 0666)
		if err != nil {
			log.Fatal().Err(err).Msg("Failed to create log file log.txt")
		}
		defer logFile.Close()
		multiLogWriter := io.MultiWriter(os.Stdout, logFile)

		// set some default fields added to all logs
		log.Logger = zerolog.New(multiLogWriter).With().
			Timestamp().
			Str("app", "estafette-ci-builder").
			Str("version", version).
			Str("jobName", jobName).
			Str("gitName", gitName).
			Str("gitBranch", gitBranch).
			Str("gitRevision", gitRevision).
			Logger()

		stdlog.SetFlags(0)
		stdlog.SetOutput(log.Logger)

		// log startup message
		log.Info().
			Str("branch", branch).
			Str("revision", revision).
			Str("buildDate", buildDate).
			Str("goVersion", goVersion).
			Msg("Starting estafette-ci-builder...")

		// start docker daemon
		err = dockerRunner.startDockerDaemon()
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Error starting docker daemon")
		}

		// wait for docker daemon to be ready for usage
		dockerRunner.waitForDockerDaemon()

		// get current working directory
		dir, err := os.Getwd()
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Getting current working directory failed")
		}

		// set some envvars
		err = envvarHelper.setEstafetteGlobalEnvvars()
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Setting global environment variables failed")
		}

		// get manifest from envvar and unmarshal
		manifestJSON := os.Getenv("ESTAFETTE_CI_MANIFEST_JSON")
		var manifest manifest.EstafetteManifest
		json.Unmarshal([]byte(manifestJSON), &manifest)

		// initialize obfuscator
		err = obfuscator.CollectSecrets(manifest)
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Collecting secrets to obfuscate failed")
		}

		// check whether this is a regular build or a release
		stages := manifest.Stages
		if releaseID > 0 {
			// check if the release is defined
			releaseExists := false
			for _, r := range manifest.Releases {
				if r.Name == releaseName {
					releaseExists = true
					stages = r.Stages
				}
			}
			if !releaseExists {
				endOfLifeHelper.handleFatal(buildLog, err, fmt.Sprintf("Release %v does not exist", releaseName))
			}
			log.Info().Msgf("Starting release %v at version %v...", releaseName, buildVersion)
		} else {
			log.Info().Msgf("Starting build version %v...", buildVersion)
		}

		// create docker client
		_, err = dockerRunner.createDockerClient()
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Failed creating a docker client")
		}

		// get private container registries credentials
		credentialsJSON := os.Getenv("ESTAFETTE_CI_REPOSITORY_CREDENTIALS_JSON")
		if credentialsJSON != "" {
			var credentials []*contracts.ContainerRepositoryCredentialConfig
			json.Unmarshal([]byte(credentialsJSON), &credentials)
			dockerRunner.setRepositoryCredentials(credentials)
		}

		// collect estafette envvars and run stages from manifest
		log.Info().Msgf("Running %v stages", len(stages))
		estafetteEnvvars := envvarHelper.collectEstafetteEnvvars(manifest)
		globalEnvvars := envvarHelper.collectGlobalEnvvars(manifest)
		envvars := envvarHelper.overrideEnvvars(estafetteEnvvars, globalEnvvars)
		result, err := pipelineRunner.runStages(stages, dir, envvars)
		if err != nil {
			endOfLifeHelper.handleFatal(buildLog, err, "Executing stages from manifest failed")
		}

		// send result to ci-api
		log.Info().Interface("result", result).Msg("Finished running stages")
		buildLog.Steps = transformPipelineRunResultToBuildLogSteps(result)
		endOfLifeHelper.sendBuildJobLogEvent(buildLog)
		buildStatus := "succeeded"
		if result.HasErrors() {
			buildStatus = "failed"
		}
		endOfLifeHelper.sendBuildFinishedEvent(buildStatus)

		if *runAsJob {
			os.Exit(0)
		} else {
			handleExit(result)
		}

	} else {
		// Set up a simple console logger
		log.Logger = zerolog.New(zerolog.ConsoleWriter{Out: os.Stderr}).With().
			Timestamp().
			Logger()

		log.Warn().Msgf("The CI Server (\"%s\") is not recognized, exiting.", ciServer)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io"
	"io/ioutil"
	_ "net/http/pprof"
	"os"
	"path"
	"runtime"
	"strconv"

	"github.com/signalfx/golib/datapoint/dpsink"
	"github.com/signalfx/golib/errors"
	"github.com/signalfx/golib/log"
	"github.com/signalfx/golib/sfxclient"
	"github.com/signalfx/golib/timekeeper"
	"github.com/signalfx/metricproxy/config"
	"github.com/signalfx/metricproxy/logkey"
	"github.com/signalfx/metricproxy/protocol"

	"expvar"
	"fmt"
	"github.com/signalfx/golib/eventcounter"
	"github.com/signalfx/metricproxy/debug"
	"github.com/signalfx/metricproxy/dp/dpbuffered"
	"github.com/signalfx/metricproxy/protocol/demultiplexer"
	"golang.org/x/net/context"
	lumberjack "gopkg.in/natefinch/lumberjack.v2"
	"strings"
	"sync"
	"time"
)

const versionString = "0.9.0"

func writePidFile(pidFileName string) error {
	pid := os.Getpid()
	return ioutil.WriteFile(pidFileName, []byte(strconv.FormatInt(int64(pid), 10)), os.FileMode(0644))
}

type proxyFlags struct {
	configFileName string
}

type proxy struct {
	flags           proxyFlags
	listeners       []protocol.Listener
	forwarders      []protocol.Forwarder
	logger          log.Logger
	setupDoneSignal chan struct{}
	tk              timekeeper.TimeKeeper
	debugServer     *debug.Server
	stdout          io.Writer
	gomaxprocs      func(int) int
}

var mainInstance = proxy{
	tk:         timekeeper.RealTime{},
	logger:     log.DefaultLogger.CreateChild(),
	stdout:     os.Stdout,
	gomaxprocs: runtime.GOMAXPROCS,
}

func init() {
	flag.StringVar(&mainInstance.flags.configFileName, "configfile", "sf/sfdbproxy.conf", "Name of the db proxy configuration file")
}

func (p *proxy) getLogOutput(loadedConfig *config.ProxyConfig) io.Writer {
	logDir := *loadedConfig.LogDir
	if logDir == "-" {
		p.logger.Log("Sending logging to stdout")
		return p.stdout
	}
	logMaxSize := *loadedConfig.LogMaxSize
	logMaxBackups := *loadedConfig.LogMaxBackups
	lumberjackLogger := &lumberjack.Logger{
		Filename:   path.Join(logDir, "metricproxy.log"),
		MaxSize:    logMaxSize, // megabytes
		MaxBackups: logMaxBackups,
	}
	p.logger.Log(logkey.Filename, lumberjackLogger.Filename, logkey.Dir, os.TempDir(), "Logging redirect setup")
	return lumberjackLogger
}

func (p *proxy) getLogger(loadedConfig *config.ProxyConfig) log.Logger {
	out := p.getLogOutput(loadedConfig)
	useJSON := *loadedConfig.LogFormat == "json"
	if useJSON {
		return log.NewJSONLogger(out, log.DefaultErrorHandler)
	}
	return log.NewLogfmtLogger(out, log.DefaultErrorHandler)
}

func forwarderName(f *config.ForwardTo) string {
	if f.Name != nil {
		return *f.Name
	}
	return f.Type
}

func getHostname(osHostname func() (string, error)) string {
	name, err := osHostname()
	if err != nil {
		return "unknown"
	}
	return name
}

func setupForwarders(ctx context.Context, hostname string, tk timekeeper.TimeKeeper, loader *config.Loader, loadedConfig *config.ProxyConfig, logger log.Logger, scheduler *sfxclient.Scheduler) ([]protocol.Forwarder, error) {
	allForwarders := make([]protocol.Forwarder, 0, len(loadedConfig.ForwardTo))
	for idx, forwardConfig := range loadedConfig.ForwardTo {
		logCtx := log.NewContext(logger).With(logkey.Protocol, forwardConfig.Type, logkey.Direction, "forwarder")
		forwarder, err := loader.Forwarder(forwardConfig)
		if err != nil {
			logCtx.Log(log.Err, err, "unable to load config")
			return nil, err
		}
		name := forwarderName(forwardConfig)
		// Buffering -> counting -> (forwarder)
		limitedLogger := &log.RateLimitedLogger{
			EventCounter: eventcounter.New(tk.Now(), time.Second),
			Limit:        2, // Only 1 a second
			Logger:       logCtx,
			Now:          tk.Now,
		}
		count := &dpsink.Counter{
			Logger: limitedLogger,
		}
		endingSink := dpsink.FromChain(forwarder, dpsink.NextWrap(count))
		bconf := &dpbuffered.Config{
			BufferSize:         forwardConfig.BufferSize,
			MaxTotalDatapoints: forwardConfig.BufferSize,
			MaxTotalEvents:     forwardConfig.BufferSize,
			MaxDrainSize:       forwardConfig.MaxDrainSize,
			NumDrainingThreads: forwardConfig.DrainingThreads,
		}
		bf := dpbuffered.NewBufferedForwarder(ctx, bconf, endingSink, limitedLogger)
		allForwarders = append(allForwarders, bf)

		groupName := fmt.Sprintf("%s_f_%d", name, idx)

		scheduler.AddGroupedCallback(groupName, forwarder)
		scheduler.GroupedDefaultDimensions(groupName, map[string]string{
			"name":      name,
			"direction": "forwarder",
			"source":    "proxy",
			"host":      hostname,
			"type":      forwardConfig.Type,
		})
	}
	return allForwarders, nil
}

func setupListeners(tk timekeeper.TimeKeeper, hostname string, loader *config.Loader, loadedConfig *config.ProxyConfig, multiplexer dpsink.Sink, logger log.Logger, scheduler *sfxclient.Scheduler) ([]protocol.Listener, error) {
	listeners := make([]protocol.Listener, 0, len(loadedConfig.ListenFrom))
	for idx, listenConfig := range loadedConfig.ListenFrom {
		logCtx := log.NewContext(logger).With(logkey.Protocol, listenConfig.Type, logkey.Direction, "listener")
		name := func() string {
			if listenConfig.Name != nil {
				return *listenConfig.Name
			}
			return listenConfig.Type
		}()
		count := &dpsink.Counter{
			Logger: &log.RateLimitedLogger{
				EventCounter: eventcounter.New(tk.Now(), time.Second),
				Limit:        1, // Only 1 a second
				Logger:       logCtx,
				Now:          tk.Now,
			},
		}
		endingSink := dpsink.FromChain(multiplexer, dpsink.NextWrap(count))

		listener, err := loader.Listener(endingSink, listenConfig)
		if err != nil {
			logCtx.Log(log.Err, err, "unable to load config")
			return nil, err
		}
		listeners = append(listeners, listener)
		groupName := fmt.Sprintf("%s_l_%d", name, idx)
		scheduler.AddGroupedCallback(groupName, listener)
		scheduler.AddGroupedCallback(groupName, count)
		scheduler.GroupedDefaultDimensions(groupName, map[string]string{
			"name":      name,
			"direction": "listener",
			"source":    "proxy",
			"host":      hostname,
			"type":      listenConfig.Type,
		})
	}
	return listeners, nil
}

func splitSinks(forwarders []protocol.Forwarder) ([]dpsink.DSink, []dpsink.ESink) {
	dsinks := make([]dpsink.DSink, 0, len(forwarders))
	esinks := make([]dpsink.ESink, 0, len(forwarders))
	for _, f := range forwarders {
		dsinks = append(dsinks, f)
		esinks = append(esinks, f)
	}
	return dsinks, esinks
}

func (p *proxy) setupDebugServer(conf *config.ProxyConfig, logger log.Logger, scheduler *sfxclient.Scheduler) error {
	if conf.LocalDebugServer == nil {
		return nil
	}
	debugServer, err := debug.New(*conf.LocalDebugServer, p, &debug.Config{
		Logger: log.NewContext(logger).With(logkey.Protocol, "debugserver"),
	})
	if err != nil {
		return errors.Annotate(err, "cannot setup debug server")
	}
	p.debugServer = debugServer
	p.debugServer.Exp2.Exported["config"] = conf.Var()
	p.debugServer.Exp2.Exported["datapoints"] = scheduler.Var()
	p.debugServer.Exp2.Exported["goruntime"] = expvar.Func(func() interface{} {
		return runtime.Version()
	})
	p.debugServer.Exp2.Exported["proxy_version"] = expvar.Func(func() interface{} {
		return versionString
	})
	return nil
}

func setupGoMaxProcs(numProcs *int, gomaxprocs func(int) int) {
	if numProcs != nil {
		gomaxprocs(*numProcs)
	} else {
		numProcs := runtime.NumCPU()
		gomaxprocs(numProcs)
	}
}

func (p *proxy) Close() error {
	errs := make([]error, len(p.forwarders)+len(p.listeners)+1)
	for _, f := range p.forwarders {
		errs = append(errs, f.Close())
	}
	for _, l := range p.listeners {
		errs = append(errs, l.Close())
	}
	if p.debugServer != nil {
		errs = append(errs, p.debugServer.Close())
	}
	return errors.NewMultiErr(errs)
}

func (p *proxy) main(ctx context.Context) error {
	// Disable the default logger to make sure nobody else uses it
	err := p.run(ctx)
	return errors.NewMultiErr([]error{err, p.Close()})
}

func (p *proxy) run(ctx context.Context) error {
	hostname := getHostname(os.Hostname)
	p.logger.Log(logkey.ConfigFile, p.flags.configFileName, "Looking for config file")
	p.logger.Log(logkey.Env, os.Environ(), "Looking for config file")

	loadedConfig, err := config.Load(p.flags.configFileName, p.logger)
	if err != nil {
		p.logger.Log(log.Err, err, "Unable to load config")
		return err
	}
	p.logger = log.NewContext(p.getLogger(loadedConfig)).With(logkey.Time, log.DefaultTimestamp, logkey.Caller, log.DefaultCaller)
	logger := p.logger
	log.DefaultLogger.Set(p.logger)
	defer func() {
		log.DefaultLogger.Set(log.Discard)
	}()

	pidFilename := *loadedConfig.PidFilename
	if err := writePidFile(pidFilename); err != nil {
		logger.Log(log.Err, err, logkey.Filename, pidFilename, "cannot store pid in pid file")
	}
	defer os.Remove(pidFilename)
	logger.Log(logkey.Config, loadedConfig, logkey.Env, strings.Join(os.Environ(), "-"), "config loaded")

	setupGoMaxProcs(loadedConfig.NumProcs, p.gomaxprocs)

	loader := config.NewLoader(ctx, logger, versionString)
	scheduler := sfxclient.NewScheduler()
	scheduler.AddCallback(sfxclient.GoMetricsSource)
	scheduler.DefaultDimensions(map[string]string{
		"source": "proxy",
		"host":   hostname,
	})

	forwarders, err := setupForwarders(ctx, hostname, p.tk, loader, loadedConfig, logger, scheduler)
	if err != nil {
		return errors.Annotate(err, "unable to setup forwarders")
	}
	p.forwarders = forwarders
	dpSinks, eSinks := splitSinks(forwarders)

	multiplexer := &demultiplexer.Demultiplexer{
		DatapointSinks: dpSinks,
		EventSinks:     eSinks,
	}

	listeners, err := setupListeners(p.tk, hostname, loader, loadedConfig, multiplexer, logger, scheduler)
	if err != nil {
		return errors.Annotate(err, "cannot setup listeners from configuration")
	}
	p.listeners = listeners

	// We still want to schedule stat collection so people can debug the server if they want
	scheduler.Sink = dpsink.Discard
	scheduler.ReportingDelayNs = (time.Second * 30).Nanoseconds()
	wg := sync.WaitGroup{}
	finisehdContext, cancelFunc := context.WithCancel(ctx)
	defer cancelFunc()
	if loadedConfig.StatsDelayDuration != nil && *loadedConfig.StatsDelayDuration != 0 {
		scheduler.Sink = multiplexer
		scheduler.ReportingDelayNs = loadedConfig.StatsDelayDuration.Nanoseconds()
	} else {
		logger.Log("skipping stat keeping")
	}

	// Schedule datapoint collection to a Discard sink so we can get the stats in Expvar()
	wg.Add(1)
	go func() {
		err := scheduler.Schedule(finisehdContext)
		logger.Log(log.Err, err, logkey.Struct, "scheduler", "Schedule finished")
		wg.Done()
	}()
	if err := p.setupDebugServer(loadedConfig, logger, scheduler); err != nil {
		return err
	}
	logger.Log("Setup done.  Blocking!")
	if p.setupDoneSignal != nil {
		close(p.setupDoneSignal)
	}
	<-ctx.Done()
	cancelFunc()
	wg.Wait()
	return nil
}

var flagParse = flag.Parse

func main() {
	flagParse()
	mainInstance.main(context.Background())
}
<commit_msg>Include buffering stats<commit_after>package main

import (
	"flag"
	"io"
	"io/ioutil"
	_ "net/http/pprof"
	"os"
	"path"
	"runtime"
	"strconv"

	"github.com/signalfx/golib/datapoint/dpsink"
	"github.com/signalfx/golib/errors"
	"github.com/signalfx/golib/log"
	"github.com/signalfx/golib/sfxclient"
	"github.com/signalfx/golib/timekeeper"
	"github.com/signalfx/metricproxy/config"
	"github.com/signalfx/metricproxy/logkey"
	"github.com/signalfx/metricproxy/protocol"

	"expvar"
	"fmt"
	"github.com/signalfx/golib/eventcounter"
	"github.com/signalfx/metricproxy/debug"
	"github.com/signalfx/metricproxy/dp/dpbuffered"
	"github.com/signalfx/metricproxy/protocol/demultiplexer"
	"golang.org/x/net/context"
	lumberjack "gopkg.in/natefinch/lumberjack.v2"
	"strings"
	"sync"
	"time"
)

const versionString = "0.9.0"

func writePidFile(pidFileName string) error {
	pid := os.Getpid()
	return ioutil.WriteFile(pidFileName, []byte(strconv.FormatInt(int64(pid), 10)), os.FileMode(0644))
}

type proxyFlags struct {
	configFileName string
}

type proxy struct {
	flags           proxyFlags
	listeners       []protocol.Listener
	forwarders      []protocol.Forwarder
	logger          log.Logger
	setupDoneSignal chan struct{}
	tk              timekeeper.TimeKeeper
	debugServer     *debug.Server
	stdout          io.Writer
	gomaxprocs      func(int) int
}

var mainInstance = proxy{
	tk:         timekeeper.RealTime{},
	logger:     log.DefaultLogger.CreateChild(),
	stdout:     os.Stdout,
	gomaxprocs: runtime.GOMAXPROCS,
}

func init() {
	flag.StringVar(&mainInstance.flags.configFileName, "configfile", "sf/sfdbproxy.conf", "Name of the db proxy configuration file")
}

func (p *proxy) getLogOutput(loadedConfig *config.ProxyConfig) io.Writer {
	logDir := *loadedConfig.LogDir
	if logDir == "-" {
		p.logger.Log("Sending logging to stdout")
		return p.stdout
	}
	logMaxSize := *loadedConfig.LogMaxSize
	logMaxBackups := *loadedConfig.LogMaxBackups
	lumberjackLogger := &lumberjack.Logger{
		Filename:   path.Join(logDir, "metricproxy.log"),
		MaxSize:    logMaxSize, // megabytes
		MaxBackups: logMaxBackups,
	}
	p.logger.Log(logkey.Filename, lumberjackLogger.Filename, logkey.Dir, os.TempDir(), "Logging redirect setup")
	return lumberjackLogger
}

func (p *proxy) getLogger(loadedConfig *config.ProxyConfig) log.Logger {
	out := p.getLogOutput(loadedConfig)
	useJSON := *loadedConfig.LogFormat == "json"
	if useJSON {
		return log.NewJSONLogger(out, log.DefaultErrorHandler)
	}
	return log.NewLogfmtLogger(out, log.DefaultErrorHandler)
}

func forwarderName(f *config.ForwardTo) string {
	if f.Name != nil {
		return *f.Name
	}
	return f.Type
}

func getHostname(osHostname func() (string, error)) string {
	name, err := osHostname()
	if err != nil {
		return "unknown"
	}
	return name
}

func setupForwarders(ctx context.Context, hostname string, tk timekeeper.TimeKeeper, loader *config.Loader, loadedConfig *config.ProxyConfig, logger log.Logger, scheduler *sfxclient.Scheduler) ([]protocol.Forwarder, error) {
	allForwarders := make([]protocol.Forwarder, 0, len(loadedConfig.ForwardTo))
	for idx, forwardConfig := range loadedConfig.ForwardTo {
		logCtx := log.NewContext(logger).With(logkey.Protocol, forwardConfig.Type, logkey.Direction, "forwarder")
		forwarder, err := loader.Forwarder(forwardConfig)
		if err != nil {
			logCtx.Log(log.Err, err, "unable to load config")
			return nil, err
		}
		name := forwarderName(forwardConfig)
		// Buffering -> counting -> (forwarder)
		limitedLogger := &log.RateLimitedLogger{
			EventCounter: eventcounter.New(tk.Now(), time.Second),
			Limit:        2, // Only 1 a second
			Logger:       logCtx,
			Now:          tk.Now,
		}
		count := &dpsink.Counter{
			Logger: limitedLogger,
		}
		endingSink := dpsink.FromChain(forwarder, dpsink.NextWrap(count))
		bconf := &dpbuffered.Config{
			BufferSize:         forwardConfig.BufferSize,
			MaxTotalDatapoints: forwardConfig.BufferSize,
			MaxTotalEvents:     forwardConfig.BufferSize,
			MaxDrainSize:       forwardConfig.MaxDrainSize,
			NumDrainingThreads: forwardConfig.DrainingThreads,
		}
		bf := dpbuffered.NewBufferedForwarder(ctx, bconf, endingSink, limitedLogger)
		allForwarders = append(allForwarders, bf)

		groupName := fmt.Sprintf("%s_f_%d", name, idx)

		scheduler.AddGroupedCallback(groupName, forwarder)
		scheduler.AddGroupedCallback(groupName, bf)
		scheduler.GroupedDefaultDimensions(groupName, map[string]string{
			"name":      name,
			"direction": "forwarder",
			"source":    "proxy",
			"host":      hostname,
			"type":      forwardConfig.Type,
		})
	}
	return allForwarders, nil
}

func setupListeners(tk timekeeper.TimeKeeper, hostname string, loader *config.Loader, loadedConfig *config.ProxyConfig, multiplexer dpsink.Sink, logger log.Logger, scheduler *sfxclient.Scheduler) ([]protocol.Listener, error) {
	listeners := make([]protocol.Listener, 0, len(loadedConfig.ListenFrom))
	for idx, listenConfig := range loadedConfig.ListenFrom {
		logCtx := log.NewContext(logger).With(logkey.Protocol, listenConfig.Type, logkey.Direction, "listener")
		name := func() string {
			if listenConfig.Name != nil {
				return *listenConfig.Name
			}
			return listenConfig.Type
		}()
		count := &dpsink.Counter{
			Logger: &log.RateLimitedLogger{
				EventCounter: eventcounter.New(tk.Now(), time.Second),
				Limit:        1, // Only 1 a second
				Logger:       logCtx,
				Now:          tk.Now,
			},
		}
		endingSink := dpsink.FromChain(multiplexer, dpsink.NextWrap(count))

		listener, err := loader.Listener(endingSink, listenConfig)
		if err != nil {
			logCtx.Log(log.Err, err, "unable to load config")
			return nil, err
		}
		listeners = append(listeners, listener)
		groupName := fmt.Sprintf("%s_l_%d", name, idx)
		scheduler.AddGroupedCallback(groupName, listener)
		scheduler.AddGroupedCallback(groupName, count)
		scheduler.GroupedDefaultDimensions(groupName, map[string]string{
			"name":      name,
			"direction": "listener",
			"source":    "proxy",
			"host":      hostname,
			"type":      listenConfig.Type,
		})
	}
	return listeners, nil
}

func splitSinks(forwarders []protocol.Forwarder) ([]dpsink.DSink, []dpsink.ESink) {
	dsinks := make([]dpsink.DSink, 0, len(forwarders))
	esinks := make([]dpsink.ESink, 0, len(forwarders))
	for _, f := range forwarders {
		dsinks = append(dsinks, f)
		esinks = append(esinks, f)
	}
	return dsinks, esinks
}

func (p *proxy) setupDebugServer(conf *config.ProxyConfig, logger log.Logger, scheduler *sfxclient.Scheduler) error {
	if conf.LocalDebugServer == nil {
		return nil
	}
	debugServer, err := debug.New(*conf.LocalDebugServer, p, &debug.Config{
		Logger: log.NewContext(logger).With(logkey.Protocol, "debugserver"),
	})
	if err != nil {
		return errors.Annotate(err, "cannot setup debug server")
	}
	p.debugServer = debugServer
	p.debugServer.Exp2.Exported["config"] = conf.Var()
	p.debugServer.Exp2.Exported["datapoints"] = scheduler.Var()
	p.debugServer.Exp2.Exported["goruntime"] = expvar.Func(func() interface{} {
		return runtime.Version()
	})
	p.debugServer.Exp2.Exported["proxy_version"] = expvar.Func(func() interface{} {
		return versionString
	})
	return nil
}

func setupGoMaxProcs(numProcs *int, gomaxprocs func(int) int) {
	if numProcs != nil {
		gomaxprocs(*numProcs)
	} else {
		numProcs := runtime.NumCPU()
		gomaxprocs(numProcs)
	}
}

func (p *proxy) Close() error {
	errs := make([]error, len(p.forwarders)+len(p.listeners)+1)
	for _, f := range p.forwarders {
		errs = append(errs, f.Close())
	}
	for _, l := range p.listeners {
		errs = append(errs, l.Close())
	}
	if p.debugServer != nil {
		errs = append(errs, p.debugServer.Close())
	}
	return errors.NewMultiErr(errs)
}

func (p *proxy) main(ctx context.Context) error {
	// Disable the default logger to make sure nobody else uses it
	err := p.run(ctx)
	return errors.NewMultiErr([]error{err, p.Close()})
}

func (p *proxy) run(ctx context.Context) error {
	hostname := getHostname(os.Hostname)
	p.logger.Log(logkey.ConfigFile, p.flags.configFileName, "Looking for config file")
	p.logger.Log(logkey.Env, os.Environ(), "Looking for config file")

	loadedConfig, err := config.Load(p.flags.configFileName, p.logger)
	if err != nil {
		p.logger.Log(log.Err, err, "Unable to load config")
		return err
	}
	p.logger = log.NewContext(p.getLogger(loadedConfig)).With(logkey.Time, log.DefaultTimestamp, logkey.Caller, log.DefaultCaller)
	logger := p.logger
	log.DefaultLogger.Set(p.logger)
	defer func() {
		log.DefaultLogger.Set(log.Discard)
	}()

	pidFilename := *loadedConfig.PidFilename
	if err := writePidFile(pidFilename); err != nil {
		logger.Log(log.Err, err, logkey.Filename, pidFilename, "cannot store pid in pid file")
	}
	defer os.Remove(pidFilename)
	logger.Log(logkey.Config, loadedConfig, logkey.Env, strings.Join(os.Environ(), "-"), "config loaded")

	setupGoMaxProcs(loadedConfig.NumProcs, p.gomaxprocs)

	loader := config.NewLoader(ctx, logger, versionString)
	scheduler := sfxclient.NewScheduler()
	scheduler.AddCallback(sfxclient.GoMetricsSource)
	scheduler.DefaultDimensions(map[string]string{
		"source": "proxy",
		"host":   hostname,
	})

	forwarders, err := setupForwarders(ctx, hostname, p.tk, loader, loadedConfig, logger, scheduler)
	if err != nil {
		return errors.Annotate(err, "unable to setup forwarders")
	}
	p.forwarders = forwarders
	dpSinks, eSinks := splitSinks(forwarders)

	multiplexer := &demultiplexer.Demultiplexer{
		DatapointSinks: dpSinks,
		EventSinks:     eSinks,
	}

	listeners, err := setupListeners(p.tk, hostname, loader, loadedConfig, multiplexer, logger, scheduler)
	if err != nil {
		return errors.Annotate(err, "cannot setup listeners from configuration")
	}
	p.listeners = listeners

	// We still want to schedule stat collection so people can debug the server if they want
	scheduler.Sink = dpsink.Discard
	scheduler.ReportingDelayNs = (time.Second * 30).Nanoseconds()
	wg := sync.WaitGroup{}
	finisehdContext, cancelFunc := context.WithCancel(ctx)
	defer cancelFunc()
	if loadedConfig.StatsDelayDuration != nil && *loadedConfig.StatsDelayDuration != 0 {
		scheduler.Sink = multiplexer
		scheduler.ReportingDelayNs = loadedConfig.StatsDelayDuration.Nanoseconds()
	} else {
		logger.Log("skipping stat keeping")
	}

	// Schedule datapoint collection to a Discard sink so we can get the stats in Expvar()
	wg.Add(1)
	go func() {
		err := scheduler.Schedule(finisehdContext)
		logger.Log(log.Err, err, logkey.Struct, "scheduler", "Schedule finished")
		wg.Done()
	}()
	if err := p.setupDebugServer(loadedConfig, logger, scheduler); err != nil {
		return err
	}
	logger.Log("Setup done.  Blocking!")
	if p.setupDoneSignal != nil {
		close(p.setupDoneSignal)
	}
	<-ctx.Done()
	cancelFunc()
	wg.Wait()
	return nil
}

var flagParse = flag.Parse

func main() {
	flagParse()
	mainInstance.main(context.Background())
}
<|endoftext|>


<commit_before>package main

import (
    "html/template"
    "log"
    "net/http"
    "crypto/rand"
    "io"
    "io/ioutil"
    "os"
    "github.com/gorilla/sessions"
    "github.com/gorilla/securecookie"
    "github.com/gorilla/websocket"
//    "github.com/gorilla/context"
    "database/sql"
)

const (
    ADDR string = ":9000"
)

var store = sessions.NewCookieStore([]byte(
                                    securecookie.GenerateRandomKey(256)))
var freeRooms = make(map[string]*room)
var allRooms = make(map[string]*room)
var roomsCount int
var roomCapacity int

// Initializing logging objects
var (
    Trace   *log.Logger
    Info    *log.Logger
    Warning *log.Logger
    Error   *log.Logger
)

type appContext struct {
  db *sql.DB
}

func Init(
    traceHandle io.Writer,
    infoHandle io.Writer,
    warningHandle io.Writer,
    errorHandle io.Writer) {

    Trace = log.New(traceHandle,
        "TRACE: ",
        log.Ldate|log.Ltime|log.Lshortfile)

    Info = log.New(infoHandle,
        "INFO: ",
        log.Ldate|log.Ltime|log.Lshortfile)

    Warning = log.New(warningHandle,
        "WARNING: ",
        log.Ldate|log.Ltime|log.Lshortfile)

    Error = log.New(errorHandle,
        "ERROR: ",
        log.Ldate|log.Ltime|log.Lshortfile)
}

/*
    Structures used
*/

type room struct {
    name string

    // Registered connections.
    playerConns map[*playerConn]bool

    // Update state for all conn.
    updateAll chan bool

    // Register requests from the connections.
    joinChannel chan *playerConn

    // Unregister requests from connections.
    leaveChannel chan *playerConn

    //Channel for transferring questions
    questionChannel chan Question

    questions []Question
}

type Player struct {
    name  string
}

type playerConn struct {
    ws *websocket.Conn
    player *Player
    room *room
    question string
}

type Message struct{
    message string
}

type Command struct {
    Cmd string
    Payload string
    Sender string
}

type Question struct {
    Question string
    Answer string
    PlayerName string
}
/*
    Utility function generates random id.
*/

func RandString(n int) string {
    const alphanum = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
    var bytes = make([]byte, n)
    rand.Read(bytes)
    for i, b := range bytes {
        bytes[i] = alphanum[b % byte(len(alphanum))]
    }
    return string(bytes)
}

/*
    Structures methods
*/

func (pc *playerConn) sendState() {
    go func() {
        msg := "msg"
        Info.Println("Send message")
        err := pc.ws.WriteMessage(websocket.TextMessage, []byte(msg))
        if err != nil {
            Info.Println("Leave due to error")
            pc.room.leaveChannel <- pc
            pc.ws.Close()
        }
    }()
}

func (r *room) updateAllPlayers() {
    for c := range r.playerConns {
        c.sendState()
    }
}

// Run the room in goroutine
func (r *room) run() {
    Info.Println("Room is running")
    // make buffer fo questions
    questions := make([]Question, roomCapacity, roomCapacity * 2)

    for {
        select {
            case c :=<-r.joinChannel:
                r.playerConns[c] = true
                // r.updateAllPlayers()
                // if room is full - delete from freeRooms
                Info.Println("Join channel")
                Info.Println(len(r.playerConns))
                Info.Println(roomCapacity)

                if len(r.playerConns) == roomCapacity {
                    Info.Println("Room ", r.name, " is full")
                    for pc, _ := range r.playerConns {
                        Info.Println("Send acknowledgement to player ", pc.player.name)

                        msg := "full"
                        cmd := &Command{
                            Cmd: msg,
                            Payload: "",
                        }

                        err := pc.ws.WriteJSON(cmd)
                        if err != nil {
                            Error.Println(err)
                            pc.room.leaveChannel <- pc
                            pc.ws.Close()
                        }
                    }
                    Info.Println("Delete room ", r.name," from free rooms")
                    delete(freeRooms, r.name)
                }
            case questionReceived := <-r.questionChannel:
                Info.Println("Question received ", questionReceived)
                r.questions = append(r.questions, questionReceived)

                if len(r.questions) == roomCapacity - 1 {
                    // Send all questions over all players
                    Info.Println("Notify all players")
                    for _, q := range r.questions {
                        for pc, _ := range r.playerConns {
                            if pc.player.name != questionReceived.PlayerName {
                                cmd := &Command{
                                    Cmd: "question",
                                    Payload: q.Question,
                                    Sender: q.PlayerName,
                                }
                                pc.ws.WriteJSON(cmd)
                            }
                        }
                    }
                    // clear questions buffer for room
                    questions = questions[0:0]
                }
            case c := <-r.leaveChannel:
                Info.Println("Leave channel")
                r.updateAllPlayers()
                delete(r.playerConns, c)
                if len(r.playerConns) == 0 {
                    goto Exit
                }
            case <-r.updateAll:
                Info.Println("Update all")
                r.updateAllPlayers()
        }
        Info.Println("Room ", r.name," is waiting for new events")
    }

    Exit:

    // delete room
    delete(allRooms, r.name)
    delete(freeRooms, r.name)
    roomsCount -= 1
    Info.Println("Room closed:", r.name)
}

func executeCommand(pc *playerConn, command string) {
        Info.Println("Command ", command, " has been executed")
        if command == "question" {
            var question Question
            var err interface{}
            err = pc.ws.ReadJSON(&question)

            if err != nil {
                Info.Println("Error, websocket will be closed")
                pc.ws.Close()
                return
            }
            Info.Println("Question was red ", question)
            // Send questions over all player
            question.PlayerName = pc.player.name
            Info.Println("Question is sent to room ", pc.room.name," question channel")
            Info.Println("Question channel length : ", len(pc.room.questionChannel))
            pc.room.questionChannel <- question
        }
}

func (pc *playerConn) receiver() {
    for {
        _, command, err := pc.ws.ReadMessage()
        if err != nil {
            break
        }
        Info.Println("PC with player name ", pc.player.name,
                    " execute command: ", string(command))
        // execute a command
        // update all conn
        executeCommand(pc, string(command))
        pc.room.updateAll <- true
    }
    pc.room.leaveChannel <- pc
    pc.ws.Close()
}

func newRoom() *room {
    name := RandString(16)
    room := &room{
        name:        name,
        playerConns: make(map[*playerConn]bool),
        updateAll:   make(chan bool),
        joinChannel:        make(chan *playerConn),
        leaveChannel:       make(chan *playerConn),
        questions: make([]Question, roomCapacity, roomCapacity * 2),
    }

    allRooms[name] = room
    freeRooms[name] = room

    // run room
    go room.run()

    roomsCount += 1
    return room
}

func newPlayer(username string) *Player{
    return &Player{
        name:        username,
    }
}

func newPlayerConn(ws *websocket.Conn, player *Player, room *room) *playerConn {
    pc := &playerConn{ws, player, room, ""}
    go pc.receiver()
    return pc
}

// Function that gets request handler and covers it with authentication.
//func  (c *appContext) authHandler(next http.Handler) http.Handler {
//    fn := func(w http.ResponseWriter, r *http.Request) {
//        authToken := r.Header.Get("Authorization")
//        Info.Println(authToken)
//        user, err := "username", nil //getUser(authToken)
//        //user, err := getUser(c.db, authToken)
//
//        if err != nil {
//            Info.Println("Redirecting to login")
//            http.Redirect(w, r, "/login", 302)
//            return
//        }
//
//        context.Set(r, "user", user)
//        next.ServeHTTP(w, r)
//    }
//
//    return http.HandlerFunc(fn)
//}

/*
    HTTP handlers
*/

func homeHandler(c http.ResponseWriter, r *http.Request) {
    Info.Println("Index handler")
    var homeTempl = template.Must(template.ParseFiles("templates/home.html"))
    session, _ := store.Get(r, "dating")
    username := session.Values["username"]

    if username == nil {
        Info.Println("Redirecting to login")
        http.Redirect(c, r, "/login", 302)
        return
    }

    user_name := username.([]string)[0]
    Info.Println(user_name)

    data := struct {
        Host       string
        RoomsCount int
        Username interface{}
        rooms map[string]*room
    }{r.Host, 0, user_name, freeRooms}
    homeTempl.Execute(c, data)
}

func loginHandler(c http.ResponseWriter,r *http.Request) {
    Info.Println("Login handler")
    var loginTempl = template.Must(template.ParseFiles("templates/login.html"))

    if r.Method == "GET" {
        log.Println("Login method GET")
        data := struct {
            Error      string
        }{""}

        loginTempl.Execute(c, data)
    } else {
        session, _ := store.Get(r, "dating")

        r.ParseForm()
        Info.Println("Login POST")
        username := r.Form["username"]

        Info.Println("username ", username)
        if username == nil {
            data := struct {
                Host       string
                Error      string
            }{r.Host + "/login", "Please, enter username"}

            loginTempl.Execute(c, data)
            return
        }
        Info.Println("session ", session)
        session.Values["username"] = username
        session.Save(r, c)
        Info.Println("Redirecting to home page")
        http.Redirect(c, r, "/", 302)
    }
}

func joinOrCreateRoom(c http.ResponseWriter,r *http.Request) {
    ws, err := websocket.Upgrade(c, r, nil, 1024, 1024)
    Info.Println("Join room")
    if _, ok := err.(websocket.HandshakeError); ok {
        http.Error(c, "Not a websocket handshake", 400)
        return
    } else if err != nil {
        return
    }

    /*
        TODO: should be extracted to separate function.
    */
    session, _ := store.Get(r, "dating")
    user_data := session.Values["username"]

    var freeRoom *room
    if len(freeRooms) != 0 {
        for _, room := range freeRooms{
            freeRoom = room
            break
        }
    } else {
        Info.Println("New room!!!")
        freeRoom = newRoom()
    }
    Info.Println("Username ", user_data.([]string)[0],
                " has connected")
    username := user_data.([]string)[0]
    player := newPlayer(username)
    pConn := newPlayerConn(ws, player, freeRoom)
    // Join Player to room
    freeRoom.joinChannel <- pConn
}

func main() {
    Init(ioutil.Discard, os.Stdout, os.Stdout, os.Stderr)

    roomCapacity = 1
    http.HandleFunc("/", homeHandler)
    http.HandleFunc("/login", loginHandler)
    http.HandleFunc("/join", joinOrCreateRoom)

    Info.Println("Start server on port", ADDR[1:])
    http.HandleFunc("/static/", func(w http.ResponseWriter, r *http.Request) {
        http.ServeFile(w, r, r.URL.Path[1:])
    })

    if err := http.ListenAndServe(ADDR, nil); err != nil {
        log.Fatal("ListenAndServe:", err)
    }
}
<commit_msg>question senging fixed<commit_after>package main

import (
    "html/template"
    "log"
    "net/http"
    "crypto/rand"
    "io"
    "io/ioutil"
    "os"
    "github.com/gorilla/sessions"
    "github.com/gorilla/securecookie"
    "github.com/gorilla/websocket"
//    "github.com/gorilla/context"
    "database/sql"
)

const (
    ADDR string = ":9000"
)

var store = sessions.NewCookieStore([]byte(
                                    securecookie.GenerateRandomKey(256)))
var freeRooms = make(map[string]*room)
var allRooms = make(map[string]*room)
var roomsCount int
var roomCapacity int

// Initializing logging objects
var (
    Trace   *log.Logger
    Info    *log.Logger
    Warning *log.Logger
    Error   *log.Logger
)

type appContext struct {
  db *sql.DB
}

func Init(
    traceHandle io.Writer,
    infoHandle io.Writer,
    warningHandle io.Writer,
    errorHandle io.Writer) {

    Trace = log.New(traceHandle,
        "TRACE: ",
        log.Ldate|log.Ltime|log.Lshortfile)

    Info = log.New(infoHandle,
        "INFO: ",
        log.Ldate|log.Ltime|log.Lshortfile)

    Warning = log.New(warningHandle,
        "WARNING: ",
        log.Ldate|log.Ltime|log.Lshortfile)

    Error = log.New(errorHandle,
        "ERROR: ",
        log.Ldate|log.Ltime|log.Lshortfile)
}

/*
    Structures used
*/

type room struct {
    name string

    // Registered connections.
    playerConns map[*playerConn]bool

    // Update state for all conn.
    updateAll chan bool

    // Register requests from the connections.
    joinChannel chan *playerConn

    // Unregister requests from connections.
    leaveChannel chan *playerConn

    //Channel for transferring questions
    questionChannel chan Question

    questions []Question
}

type Player struct {
    name  string
}

type playerConn struct {
    ws *websocket.Conn
    player *Player
    room *room
    question string
}

type Message struct{
    message string
}

type Command struct {
    Cmd string
    Payload string
    Sender string
}

type Question struct {
    Question string
    Answer string
    PlayerName string
}
/*
    Utility function generates random id.
*/

func RandString(n int) string {
    const alphanum = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
    var bytes = make([]byte, n)
    rand.Read(bytes)
    for i, b := range bytes {
        bytes[i] = alphanum[b % byte(len(alphanum))]
    }
    return string(bytes)
}

/*
    Structures methods
*/

func (pc *playerConn) sendState() {
    go func() {
        msg := "msg"
        Info.Println("Send message")
        err := pc.ws.WriteMessage(websocket.TextMessage, []byte(msg))
        if err != nil {
            Info.Println("Leave due to error")
            pc.room.leaveChannel <- pc
            pc.ws.Close()
        }
    }()
}

func (r *room) updateAllPlayers() {
    for c := range r.playerConns {
        c.sendState()
    }
}

// Run the room in goroutine
func (r *room) run() {
    Info.Println("Room is running")
    // make buffer fo questions
    questions := make([]Question, roomCapacity, roomCapacity * 2)

    for {
        select {
            case c :=<-r.joinChannel:
                r.playerConns[c] = true
                // r.updateAllPlayers()
                // if room is full - delete from freeRooms
                Info.Println("Join channel")
                Info.Println(len(r.playerConns))
                Info.Println(roomCapacity)

                if len(r.playerConns) == roomCapacity {
                    Info.Println("Room ", r.name, " is full")
                    for pc, _ := range r.playerConns {
                        Info.Println("Send acknowledgement to player ", pc.player.name)

                        msg := "full"
                        cmd := &Command{
                            Cmd: msg,
                            Payload: "",
                        }

                        err := pc.ws.WriteJSON(cmd)
                        if err != nil {
                            Error.Println(err)
                            pc.room.leaveChannel <- pc
                            pc.ws.Close()
                        }
                    }
                    Info.Println("Delete room ", r.name," from free rooms")
                    delete(freeRooms, r.name)
                }
            case questionReceived := <-r.questionChannel:
                Info.Println("Question received ", questionReceived)
                r.questions = append(r.questions, questionReceived)

                if len(r.questions) == roomCapacity - 1 {
                    // Send all questions over all players
                    Info.Println("Notify all players")
                    for _, q := range r.questions {
                        for pc, _ := range r.playerConns {
                            if pc.player.name != questionReceived.PlayerName {
                                cmd := &Command{
                                    Cmd: "question",
                                    Payload: q.Question,
                                    Sender: q.PlayerName,
                                }
                                pc.ws.WriteJSON(cmd)
                            }
                        }
                    }
                    // clear questions buffer for room
                    questions = questions[0:0]
                }
            case c := <-r.leaveChannel:
                Info.Println("Leave channel")
                r.updateAllPlayers()
                delete(r.playerConns, c)
                if len(r.playerConns) == 0 {
                    goto Exit
                }
            case <-r.updateAll:
                Info.Println("Update all")
                r.updateAllPlayers()
        }
        Info.Println("Room ", r.name," is waiting for new events")
    }

    Exit:

    // delete room
    delete(allRooms, r.name)
    delete(freeRooms, r.name)
    roomsCount -= 1
    Info.Println("Room closed:", r.name)
}

func executeCommand(pc *playerConn, command string) {
        Info.Println("Command ", command, " has been executed")
        if command == "question" {
            var question Question
            var err interface{}
            err = pc.ws.ReadJSON(&question)

            if err != nil {
                Info.Println("Error, websocket will be closed")
                pc.ws.Close()
                return
            }
            Info.Println("Question was red ", question)
            // Send questions over all player
            question.PlayerName = pc.player.name
            Info.Println("Question is sent to room ", pc.room.name," question channel")
            Info.Println("Question channel length : ", len(pc.room.questionChannel))
            pc.room.questionChannel <- question
        }
}

func (pc *playerConn) receiver() {
    for {
        _, command, err := pc.ws.ReadMessage()
        if err != nil {
            break
        }
        Info.Println("PC with player name ", pc.player.name,
                    " execute command: ", string(command))
        // execute a command
        // update all conn
        executeCommand(pc, string(command))
        pc.room.updateAll <- true
    }
    pc.room.leaveChannel <- pc
    pc.ws.Close()
}

func newRoom() *room {
    name := RandString(16)
    room := &room{
        name:        name,
        playerConns: make(map[*playerConn]bool),
        updateAll:   make(chan bool),
        joinChannel:        make(chan *playerConn),
        leaveChannel:       make(chan *playerConn),
        questionChannel: make(chan Question),
        questions: make([]Question, roomCapacity, roomCapacity * 2),
    }

    allRooms[name] = room
    freeRooms[name] = room

    // run room
    go room.run()

    roomsCount += 1
    return room
}

func newPlayer(username string) *Player{
    return &Player{
        name:        username,
    }
}

func newPlayerConn(ws *websocket.Conn, player *Player, room *room) *playerConn {
    pc := &playerConn{ws, player, room, ""}
    go pc.receiver()
    return pc
}

// Function that gets request handler and covers it with authentication.
//func  (c *appContext) authHandler(next http.Handler) http.Handler {
//    fn := func(w http.ResponseWriter, r *http.Request) {
//        authToken := r.Header.Get("Authorization")
//        Info.Println(authToken)
//        user, err := "username", nil //getUser(authToken)
//        //user, err := getUser(c.db, authToken)
//
//        if err != nil {
//            Info.Println("Redirecting to login")
//            http.Redirect(w, r, "/login", 302)
//            return
//        }
//
//        context.Set(r, "user", user)
//        next.ServeHTTP(w, r)
//    }
//
//    return http.HandlerFunc(fn)
//}

/*
    HTTP handlers
*/

func homeHandler(c http.ResponseWriter, r *http.Request) {
    Info.Println("Index handler")
    var homeTempl = template.Must(template.ParseFiles("templates/home.html"))
    session, _ := store.Get(r, "dating")
    username := session.Values["username"]

    if username == nil {
        Info.Println("Redirecting to login")
        http.Redirect(c, r, "/login", 302)
        return
    }

    user_name := username.([]string)[0]
    Info.Println(user_name)

    data := struct {
        Host       string
        RoomsCount int
        Username interface{}
        rooms map[string]*room
    }{r.Host, 0, user_name, freeRooms}
    homeTempl.Execute(c, data)
}

func loginHandler(c http.ResponseWriter,r *http.Request) {
    Info.Println("Login handler")
    var loginTempl = template.Must(template.ParseFiles("templates/login.html"))

    if r.Method == "GET" {
        log.Println("Login method GET")
        data := struct {
            Error      string
        }{""}

        loginTempl.Execute(c, data)
    } else {
        session, _ := store.Get(r, "dating")

        r.ParseForm()
        Info.Println("Login POST")
        username := r.Form["username"]

        Info.Println("username ", username)
        if username == nil {
            data := struct {
                Host       string
                Error      string
            }{r.Host + "/login", "Please, enter username"}

            loginTempl.Execute(c, data)
            return
        }
        Info.Println("session ", session)
        session.Values["username"] = username
        session.Save(r, c)
        Info.Println("Redirecting to home page")
        http.Redirect(c, r, "/", 302)
    }
}

func joinOrCreateRoom(c http.ResponseWriter,r *http.Request) {
    ws, err := websocket.Upgrade(c, r, nil, 1024, 1024)
    Info.Println("Join room")
    if _, ok := err.(websocket.HandshakeError); ok {
        http.Error(c, "Not a websocket handshake", 400)
        return
    } else if err != nil {
        return
    }

    /*
        TODO: should be extracted to separate function.
    */
    session, _ := store.Get(r, "dating")
    user_data := session.Values["username"]

    var freeRoom *room
    if len(freeRooms) != 0 {
        for _, room := range freeRooms{
            freeRoom = room
            break
        }
    } else {
        Info.Println("New room!!!")
        freeRoom = newRoom()
    }
    Info.Println("Username ", user_data.([]string)[0],
                " has connected")
    username := user_data.([]string)[0]
    player := newPlayer(username)
    pConn := newPlayerConn(ws, player, freeRoom)

    // Join Player to room
    freeRoom.joinChannel <- pConn
}

func main() {
    Init(ioutil.Discard, os.Stdout, os.Stdout, os.Stderr)

    roomCapacity = 1
    http.HandleFunc("/", homeHandler)
    http.HandleFunc("/login", loginHandler)
    http.HandleFunc("/join", joinOrCreateRoom)

    Info.Println("Start server on port", ADDR[1:])
    http.HandleFunc("/static/", func(w http.ResponseWriter, r *http.Request) {
        http.ServeFile(w, r, r.URL.Path[1:])
    })

    if err := http.ListenAndServe(ADDR, nil); err != nil {
        log.Fatal("ListenAndServe:", err)
    }
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"sync"

	"github.com/kori/gone/svc/teknik"
)

func main() {
	flag.Parse()
	var wg sync.WaitGroup

	// t = total files, p = path to file
	for t, p := range flag.Args() {
		wg.Add(1)
		go func(file string) {
			defer wg.Done()
			url := teknik.Upload(p)
			fmt.Println(url)
		}(p)
	}
	wg.Wait()
}
<commit_msg>Remove unused var.<commit_after>package main

import (
	"flag"
	"fmt"
	"sync"

	"github.com/kori/gone/svc/teknik"
)

func main() {
	flag.Parse()
	var wg sync.WaitGroup

	// p = path to file
	for _, p := range flag.Args() {
		wg.Add(1)
		go func(file string) {
			defer wg.Done()
			url := teknik.Upload(p)
			fmt.Println(url)
		}(p)
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"bufio"
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/nmeum/marvin/irc"
	"github.com/nmeum/marvin/modules"
	"github.com/nmeum/marvin/modules/url"
	"log"
	"net"
	"os"
	"time"
)

const (
	appName  = "marvin"
	appUsage = "[options] CHANNEL..."
)

var (
	nick = flag.String("n", "marvin", "nickname")
	name = flag.String("r", "Marvin Bot", "realname")
	host = flag.String("h", "irc.hackint.eu", "host")
	ussl = flag.Bool("s", false, "enable ssl")
	issl = flag.Bool("i", false, "ignore invalid ssl certs")
	port = flag.Int("p", 6667, "port")
)

type moduleInit func(*modules.ModuleSet)

var moduleInits = []moduleInit{
	url.Init,
}

func main() {
	flag.Parse()
	logger := log.New(os.Stderr, fmt.Sprintf("%s:", appName), 0)
	if flag.NArg() < 1 {
		logger.Fatalf("USAGE: %s %s", appName, appUsage)
	}

	conn, err := connect("tcp", fmt.Sprintf("%s:%d", *host, *port))
	if err != nil {
		logger.Fatal(err)
	}
	defer conn.Close()

	for {
		ircBot, err := setup(conn, flag.Args())
		if err != nil {
			logger.Println(err)
		}

		scanner := bufio.NewScanner(conn)
		for scanner.Scan() {
			errChan := ircBot.Handle(scanner.Text())
			for err := range errChan {
				logger.Println(err)
			}
		}

		if err := scanner.Err(); err != nil {
			logger.Println(err)
		}

		conn = reconnect(conn)
	}
}

func setup(conn net.Conn, channels []string) (client *irc.Client, err error) {
	client = irc.NewClient(&conn)
	client.CmdHook("ping", func(c *irc.Client, m irc.Message) error {
		return c.Write("PONG %s", m.Data)
	})

	client.CmdHook("001", func(c *irc.Client, m irc.Message) error {
		for _, ch := range channels {
			if err := c.Write("JOIN %s", ch); err != nil {
				return err
			}
		}

		return nil
	})

	client.Write("USER %s %s * :%s", *nick, *host, *name)
	client.Write("NICK %s", *nick)

	err = initializeModules(client)
	return
}

func initializeModules(c *irc.Client) error {
	moduleSet := modules.NewModuleSet(c)
	for _, fn := range moduleInits {
		fn(moduleSet)
	}

	return moduleSet.LoadAll()
}

func connect(network, address string) (conn net.Conn, err error) {
	if *ussl {
		config := new(tls.Config)
		if *issl {
			config.InsecureSkipVerify = true
		}

		return tls.Dial(network, address, config)
	}

	return net.Dial(network, address)
}

func reconnect(c net.Conn) net.Conn {
	addr := c.RemoteAddr()
	c.Close()

	var err error
	var conn net.Conn

	for i := 1; ; i++ {
		conn, err = connect(addr.Network(), addr.String())
		if err == nil {
			break
		}

		time.Sleep((time.Duration)(i*3) * time.Second)
	}

	return conn
}
<commit_msg>defensive programming<commit_after>// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"bufio"
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/nmeum/marvin/irc"
	"github.com/nmeum/marvin/modules"
	"github.com/nmeum/marvin/modules/url"
	"log"
	"net"
	"os"
	"time"
)

const (
	appName  = "marvin"
	appUsage = "[options] CHANNEL..."
)

var (
	nick = flag.String("n", "marvin", "nickname")
	name = flag.String("r", "Marvin Bot", "realname")
	host = flag.String("h", "irc.hackint.eu", "host")
	ussl = flag.Bool("s", false, "enable ssl")
	issl = flag.Bool("i", false, "ignore invalid ssl certs")
	port = flag.Int("p", 6667, "port")
)

type moduleInit func(*modules.ModuleSet)

var moduleInits = []moduleInit{
	url.Init,
}

func main() {
	flag.Parse()
	logger := log.New(os.Stderr, fmt.Sprintf("%s:", appName), 0)
	if flag.NArg() < 1 {
		logger.Fatalf("USAGE: %s %s", appName, appUsage)
	}

	conn, err := connect("tcp", fmt.Sprintf("%s:%d", *host, *port))
	if err != nil {
		logger.Fatal(err)
	}
	defer conn.Close()

	for {
		ircBot, err := setup(conn, flag.Args())
		if err != nil {
			logger.Println(err)
		}

		scanner := bufio.NewScanner(conn)
		for scanner.Scan() {
			data := scanner.Text()
			fmt.Println(data)

			errChan := ircBot.Handle(data)
			for err := range errChan {
				logger.Println(err)
			}
		}

		if err := scanner.Err(); err != nil {
			logger.Println(err)
		}

		conn = reconnect(conn)
	}
}

func setup(conn net.Conn, channels []string) (client *irc.Client, err error) {
	client = irc.NewClient(&conn)
	client.CmdHook("ping", func(c *irc.Client, m irc.Message) error {
		return c.Write("PONG %s", m.Data)
	})

	client.CmdHook("001", func(c *irc.Client, m irc.Message) error {
		for _, ch := range channels {
			if err := c.Write("JOIN %s", ch); err != nil {
				return err
			}
		}

		return nil
	})

	client.Write("USER %s %s * :%s", *nick, *host, *name)
	client.Write("NICK %s", *nick)

	if err = initializeModules(client); err != nil {
		return
	}

	return
}

func initializeModules(c *irc.Client) error {
	moduleSet := modules.NewModuleSet(c)
	for _, fn := range moduleInits {
		fn(moduleSet)
	}

	return moduleSet.LoadAll()
}

func connect(network, address string) (conn net.Conn, err error) {
	if *ussl {
		config := new(tls.Config)
		if *issl {
			config.InsecureSkipVerify = true
		}

		return tls.Dial(network, address, config)
	}

	return net.Dial(network, address)
}

func reconnect(c net.Conn) net.Conn {
	addr := c.RemoteAddr()
	c.Close()

	var err error
	var conn net.Conn

	for i := 1; ; i++ {
		conn, err = connect(addr.Network(), addr.String())
		if err == nil {
			break
		}

		time.Sleep((time.Duration)(i*3) * time.Second)
	}

	return conn
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"hash/crc64"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/mssola/user_agent"

	"github.com/ernado/ice"
	"github.com/gortc/sdp"
	"github.com/gortc/stun"
)

var (
	portHTTP = flag.Int("port", 3000, "http server port")
	hostHTTP = flag.String("host", "localhost", "http server host")
	portSTUN = flag.Int("port-stun", stun.DefaultPort, "UDP port")
	messages = &storage{
		data: make(map[string]*storageEntry),
	}
)

var (
	bindingRequest = stun.MessageType{
		Method: stun.MethodBinding,
		Class:  stun.ClassRequest,
	}
	bindingSuccessResponse = stun.MessageType{
		Method: stun.MethodBinding,
		Class:  stun.ClassSuccessResponse,
	}
)

func processUDPPacket(addr net.Addr, b []byte, req, res *stun.Message) error {
	if !stun.IsMessage(b) {
		log.Println("packet from", addr, "is not STUN message")
		return nil
	}
	req.Raw = b
	if err := req.Decode(); err != nil {
		return err
	}
	if req.Type != bindingRequest {
		log.Println("stun: skipping", req.Type, "from", addr)
		return nil
	}
	log.Println("stun: got", req.Type)
	res.TransactionID = req.TransactionID
	res.Type = bindingSuccessResponse
	var (
		ip   net.IP
		port int
	)
	switch a := addr.(type) {
	case *net.UDPAddr:
		ip = a.IP
		port = a.Port
	default:
		panic(fmt.Sprintf("unknown addr: %v", addr))
	}
	stun.XORMappedAddress{
		IP: ip,
		Port: port,
	}.AddTo(res)
	stun.NewSoftware("cydev.ru/sdp example").AddTo(res)
	res.WriteHeader()
	messages.add(fmt.Sprintf("%s:%d", ip, port), req)
	return nil
}

type iceServerConfiguration struct {
	URLs []string `json:"urls"`
}

type iceConfiguration struct {
	Servers []iceServerConfiguration `json:"iceServers"`
}

type storageEntry struct {
	*stun.Message
	createdAt time.Time
}

func (s storageEntry) timedOut(timeout time.Time) bool {
	return s.createdAt.Before(timeout)
}

type storage struct {
	data map[string]*storageEntry
	sync.Mutex
}

func (storage) timeout() time.Time {
	return time.Now().Add(time.Second * -60)
}

func mustClone(m *stun.Message) *stun.Message {
	b := new(stun.Message)
	if err := m.CloneTo(b); err != nil {
		panic(err)
	}
	return b
}

func (s *storage) pop(addr string) *stun.Message {
	s.Lock()
	defer s.Unlock()
	if s.data[addr] == nil {
		return nil
	}
	m := mustClone(s.data[addr].Message)
	delete(s.data, addr)
	return m
}

func (s *storage) add(addr string, m *stun.Message) {
	c := new(stun.Message)
	m.CloneTo(c)
	entry := &storageEntry{
		Message:   c,
		createdAt: time.Now(),
	}
	s.Lock()
	s.data[addr] = entry
	s.Unlock()
	log.Println("storage: added", addr)
}

func (s *storage) collect() {
	s.Lock()
	var (
		toRemove = make([]string, 0, 10)
		timeout  = s.timeout()
	)
	for addr, m := range s.data {
		if m.timedOut(timeout) {
			toRemove = append(toRemove, addr)
		}
	}
	for _, addr := range toRemove {
		delete(s.data, addr)
	}
	s.Unlock()
	if len(toRemove) > 0 {
		log.Println("storage: collected", len(toRemove))
	}
}

func (s *storage) gc() {
	ticker := time.NewTicker(time.Second * 2)
	for range ticker.C {
		s.collect()
	}
}

func main() {
	flag.Parse()
	log.SetFlags(log.Lshortfile)
	fs := http.FileServer(http.Dir("static"))
	http.Handle("/", fs)
	http.HandleFunc("/ice-configuration", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Content-type", "application/json")
		encoder := json.NewEncoder(w)
		origin := r.Header.Get("Origin")
		server := "stun:rs.cydev.ru"
		if len(origin) > 0 {
			u, err := url.Parse(origin)
			if err != nil {
				log.Printf("http: failed to parse origin %q: %s", origin, err)
			} else {
				if idx := strings.LastIndex(u.Host, ":"); idx > 0 {
					u.Host = u.Host[:idx]
				}
				server = fmt.Sprintf("stun:%s:%d", u.Host, *portSTUN)
				log.Printf("http: sending ice-server %q for origin %q", server, origin)
			}
		}
		if err := encoder.Encode(iceConfiguration{
			Servers: []iceServerConfiguration{
				{URLs: []string{server}},
			},
		}); err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintln(w, "json encode:", err)
		}
	})

	mLog, err := os.OpenFile("packets.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		log.Fatalln("Failed to open log:", err)
	}
	defer mLog.Close()
	csvLog := csv.NewWriter(mLog)

	// Custom domain support.
	for _, name := range []string{
		"stun", "turn", "sdp", "web", "ice",
	} {
		body := strings.Replace(`<!DOCTYPE html>

    <meta http-equiv="Content-Type" content="text/html; charset=utf-8"/>
    <meta name="go-import" content="gortc.io/pkg git https://github.com/gortc/sdp.git">
    <meta name="go-source"
          content="gortc.io/pkg https://github.com/gortc/pkg/ https://github.com/gortc/pkg/tree/master{/dir} https://github.com/gortc/pkg/blob/master{/dir}/{file}#L{line}">
    <meta http-equiv="refresh" content="0; url=https://godoc.org/gortc.io/pkg">

Nothing to see here; move along.
`, "pkg", name, -1)
		http.HandleFunc("/" + name, func(w http.ResponseWriter, r *http.Request) {
			defer r.Body.Close()
			fmt.Fprint(w, body)
		})
	}

	http.HandleFunc("/x/sdp", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		log.Println("http:", r.Method, "request from", r.RemoteAddr)
		if r.Method == http.MethodGet {
			http.Redirect(w, r, "/sdp/", http.StatusPermanentRedirect)
			return
		}
		s := sdp.Session{}
		data, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Println("http: ReadAll body failed:", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		if s, err = sdp.DecodeSession(data, s); err != nil {
			log.Println("http: failed to decode sdp session:", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		for k, v := range s {
			fmt.Fprintf(w, `<p class="attribute">%02d %s</p>`+"\n", k, v)
			if v.Type != sdp.TypeAttribute {
				continue
			}
			if !bytes.HasPrefix(v.Value, []byte("candidate")) {
				continue
			}
			c := new(ice.Candidate)
			fmt.Fprintln(w, `<div class="stun-message">`)
			if err = ice.ParseAttribute(v.Value, c); err != nil {
				fmt.Fprintln(w, `<p class="error">failed to parse as candidate:`, err, "
")
				fmt.Fprintln(w, `</div>`)
				continue
			}
			fmt.Fprintln(w, "
parsed as candidate:", fmt.Sprintf("%+v", c), "
")
			if c.Type != ice.CandidateServerReflexive {
				fmt.Fprintln(w, `</div>`)
				continue
			}
			addr := fmt.Sprintf("%s:%d", c.ConnectionAddress, c.Port)
			m := messages.pop(addr)
			if m == nil {
				log.Println("http: no message for", addr, "in log")
				fmt.Fprintln(w, `<p class="warning">message from candidate not found in STUN log</p>`)
				fmt.Fprintln(w, `</div>`)
				continue
			}
			fmt.Fprintln(w, `<p class="success">message found in STUN log:`, m, `</p>`)
			for _, a := range m.Attributes {
				switch a.Type {
				case stun.AttrOrigin:
					fmt.Fprintf(w, "
STUN attribute %s: %q (len=%d)
", a.Type, a.Value, a.Length)
				default:
					fmt.Fprintf(w, "
STUN attribute %s: %v (len=%d)
", a.Type, a.Value, a.Length)
				}
			}
			var (
				b64             = base64.StdEncoding.EncodeToString(m.Raw)
				messageCRC64    = crc64.Checksum(m.Raw, crc64.MakeTable(crc64.ISO))
				clipID          = fmt.Sprintf("crc64-%d", messageCRC64)
				ua              = user_agent.New(r.Header.Get("User-agent"))
				bName, bVersion = ua.Browser()
			)
			if err = csvLog.Write([]string{
				addr,
				b64,
				fmt.Sprintf("%d", messageCRC64),
				bName,
				bVersion,
				ua.OS(),
			}); err != nil {
				log.Fatalln("log: failed to write:", err)
			}
			csvLog.Flush()
			fmt.Fprintln(w, `<p>dumped: <code id="`+clipID+`">stun-decode`, b64, `
				<button class="btn" data-clipboard-target="#`+clipID+`">copy</button>

`)
			fmt.Fprintln(w, `<p>crc64: <code>`, messageCRC64, `</code></p>`)
			fmt.Fprintln(w, `</div>`)
		}
	})

	var (
		addrSTUN = fmt.Sprintf(":%d", *portSTUN)
		addrHTTP = fmt.Sprintf("%s:%d", *hostHTTP, *portHTTP)
	)
	log.Println("Listening udp", addrSTUN)
	c, err := net.ListenPacket("udp", addrSTUN)
	if err != nil {
		log.Fatalf("Failed to bind udp on %s: %s", addrSTUN, err)
	}
	defer c.Close()

	// spawning storage garbage collector
	go messages.gc()

	// spawning STUN server
	go func(conn net.PacketConn) {
		log.Println("Started STUN server on", conn.LocalAddr())
		var (
			res = new(stun.Message)
			req = new(stun.Message)
			buf = make([]byte, 1024)
		)
		for {
			// ReadFrom c to buf
			n, addr, err := c.ReadFrom(buf)
			if err != nil {
				log.Fatalln("c.ReadFrom:", err)
			}
			log.Printf("udp: got packet len(%d) from %s", n, addr)
			// processing binding request
			if err = processUDPPacket(addr, buf[:n], req, res); err != nil {
				log.Println("failed to process UDP packet:", err, "from addr", addr)
			} else {
				log.Printf("stun: parsed message %q from %s", req, addr)
				if _, err = c.WriteTo(res.Raw, addr); err != nil {
					log.Println("failed to send packet:", err)
				}
			}
			res.Reset()
			req.Reset()
		}

	}(c)
	log.Println("Listening http", addrHTTP)
	log.Fatal(http.ListenAndServe(addrHTTP, nil))
}
<commit_msg>server: fix redirect to new /x/<commit_after>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"hash/crc64"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/mssola/user_agent"

	"github.com/ernado/ice"
	"github.com/gortc/sdp"
	"github.com/gortc/stun"
)

var (
	portHTTP = flag.Int("port", 3000, "http server port")
	hostHTTP = flag.String("host", "localhost", "http server host")
	portSTUN = flag.Int("port-stun", stun.DefaultPort, "UDP port")
	messages = &storage{
		data: make(map[string]*storageEntry),
	}
)

var (
	bindingRequest = stun.MessageType{
		Method: stun.MethodBinding,
		Class:  stun.ClassRequest,
	}
	bindingSuccessResponse = stun.MessageType{
		Method: stun.MethodBinding,
		Class:  stun.ClassSuccessResponse,
	}
)

func processUDPPacket(addr net.Addr, b []byte, req, res *stun.Message) error {
	if !stun.IsMessage(b) {
		log.Println("packet from", addr, "is not STUN message")
		return nil
	}
	req.Raw = b
	if err := req.Decode(); err != nil {
		return err
	}
	if req.Type != bindingRequest {
		log.Println("stun: skipping", req.Type, "from", addr)
		return nil
	}
	log.Println("stun: got", req.Type)
	res.TransactionID = req.TransactionID
	res.Type = bindingSuccessResponse
	var (
		ip   net.IP
		port int
	)
	switch a := addr.(type) {
	case *net.UDPAddr:
		ip = a.IP
		port = a.Port
	default:
		panic(fmt.Sprintf("unknown addr: %v", addr))
	}
	stun.XORMappedAddress{
		IP: ip,
		Port: port,
	}.AddTo(res)
	stun.NewSoftware("cydev.ru/sdp example").AddTo(res)
	res.WriteHeader()
	messages.add(fmt.Sprintf("%s:%d", ip, port), req)
	return nil
}

type iceServerConfiguration struct {
	URLs []string `json:"urls"`
}

type iceConfiguration struct {
	Servers []iceServerConfiguration `json:"iceServers"`
}

type storageEntry struct {
	*stun.Message
	createdAt time.Time
}

func (s storageEntry) timedOut(timeout time.Time) bool {
	return s.createdAt.Before(timeout)
}

type storage struct {
	data map[string]*storageEntry
	sync.Mutex
}

func (storage) timeout() time.Time {
	return time.Now().Add(time.Second * -60)
}

func mustClone(m *stun.Message) *stun.Message {
	b := new(stun.Message)
	if err := m.CloneTo(b); err != nil {
		panic(err)
	}
	return b
}

func (s *storage) pop(addr string) *stun.Message {
	s.Lock()
	defer s.Unlock()
	if s.data[addr] == nil {
		return nil
	}
	m := mustClone(s.data[addr].Message)
	delete(s.data, addr)
	return m
}

func (s *storage) add(addr string, m *stun.Message) {
	c := new(stun.Message)
	m.CloneTo(c)
	entry := &storageEntry{
		Message:   c,
		createdAt: time.Now(),
	}
	s.Lock()
	s.data[addr] = entry
	s.Unlock()
	log.Println("storage: added", addr)
}

func (s *storage) collect() {
	s.Lock()
	var (
		toRemove = make([]string, 0, 10)
		timeout  = s.timeout()
	)
	for addr, m := range s.data {
		if m.timedOut(timeout) {
			toRemove = append(toRemove, addr)
		}
	}
	for _, addr := range toRemove {
		delete(s.data, addr)
	}
	s.Unlock()
	if len(toRemove) > 0 {
		log.Println("storage: collected", len(toRemove))
	}
}

func (s *storage) gc() {
	ticker := time.NewTicker(time.Second * 2)
	for range ticker.C {
		s.collect()
	}
}

func main() {
	flag.Parse()
	log.SetFlags(log.Lshortfile)
	fs := http.FileServer(http.Dir("static"))
	http.Handle("/", fs)
	http.HandleFunc("/ice-configuration", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Content-type", "application/json")
		encoder := json.NewEncoder(w)
		origin := r.Header.Get("Origin")
		server := "stun:rs.cydev.ru"
		if len(origin) > 0 {
			u, err := url.Parse(origin)
			if err != nil {
				log.Printf("http: failed to parse origin %q: %s", origin, err)
			} else {
				if idx := strings.LastIndex(u.Host, ":"); idx > 0 {
					u.Host = u.Host[:idx]
				}
				server = fmt.Sprintf("stun:%s:%d", u.Host, *portSTUN)
				log.Printf("http: sending ice-server %q for origin %q", server, origin)
			}
		}
		if err := encoder.Encode(iceConfiguration{
			Servers: []iceServerConfiguration{
				{URLs: []string{server}},
			},
		}); err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintln(w, "json encode:", err)
		}
	})

	mLog, err := os.OpenFile("packets.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		log.Fatalln("Failed to open log:", err)
	}
	defer mLog.Close()
	csvLog := csv.NewWriter(mLog)

	// Custom domain support.
	for _, name := range []string{
		"stun", "turn", "sdp", "web", "ice",
	} {
		body := strings.Replace(`<!DOCTYPE html>

    <meta http-equiv="Content-Type" content="text/html; charset=utf-8"/>
    <meta name="go-import" content="gortc.io/pkg git https://github.com/gortc/sdp.git">
    <meta name="go-source"
          content="gortc.io/pkg https://github.com/gortc/pkg/ https://github.com/gortc/pkg/tree/master{/dir} https://github.com/gortc/pkg/blob/master{/dir}/{file}#L{line}">
    <meta http-equiv="refresh" content="0; url=https://godoc.org/gortc.io/pkg">

Nothing to see here; move along.
`, "pkg", name, -1)
		http.HandleFunc("/" + name, func(w http.ResponseWriter, r *http.Request) {
			defer r.Body.Close()
			fmt.Fprint(w, body)
		})
	}

	http.HandleFunc("/x/sdp", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		log.Println("http:", r.Method, "request from", r.RemoteAddr)
		if r.Method == http.MethodGet {
			http.Redirect(w, r, "/x/sdp/", http.StatusPermanentRedirect)
			return
		}
		s := sdp.Session{}
		data, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Println("http: ReadAll body failed:", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		if s, err = sdp.DecodeSession(data, s); err != nil {
			log.Println("http: failed to decode sdp session:", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		for k, v := range s {
			fmt.Fprintf(w, `<p class="attribute">%02d %s</p>`+"\n", k, v)
			if v.Type != sdp.TypeAttribute {
				continue
			}
			if !bytes.HasPrefix(v.Value, []byte("candidate")) {
				continue
			}
			c := new(ice.Candidate)
			fmt.Fprintln(w, `<div class="stun-message">`)
			if err = ice.ParseAttribute(v.Value, c); err != nil {
				fmt.Fprintln(w, `<p class="error">failed to parse as candidate:`, err, "
")
				fmt.Fprintln(w, `</div>`)
				continue
			}
			fmt.Fprintln(w, "
parsed as candidate:", fmt.Sprintf("%+v", c), "
")
			if c.Type != ice.CandidateServerReflexive {
				fmt.Fprintln(w, `</div>`)
				continue
			}
			addr := fmt.Sprintf("%s:%d", c.ConnectionAddress, c.Port)
			m := messages.pop(addr)
			if m == nil {
				log.Println("http: no message for", addr, "in log")
				fmt.Fprintln(w, `<p class="warning">message from candidate not found in STUN log</p>`)
				fmt.Fprintln(w, `</div>`)
				continue
			}
			fmt.Fprintln(w, `<p class="success">message found in STUN log:`, m, `</p>`)
			for _, a := range m.Attributes {
				switch a.Type {
				case stun.AttrOrigin:
					fmt.Fprintf(w, "
STUN attribute %s: %q (len=%d)
", a.Type, a.Value, a.Length)
				default:
					fmt.Fprintf(w, "
STUN attribute %s: %v (len=%d)
", a.Type, a.Value, a.Length)
				}
			}
			var (
				b64             = base64.StdEncoding.EncodeToString(m.Raw)
				messageCRC64    = crc64.Checksum(m.Raw, crc64.MakeTable(crc64.ISO))
				clipID          = fmt.Sprintf("crc64-%d", messageCRC64)
				ua              = user_agent.New(r.Header.Get("User-agent"))
				bName, bVersion = ua.Browser()
			)
			if err = csvLog.Write([]string{
				addr,
				b64,
				fmt.Sprintf("%d", messageCRC64),
				bName,
				bVersion,
				ua.OS(),
			}); err != nil {
				log.Fatalln("log: failed to write:", err)
			}
			csvLog.Flush()
			fmt.Fprintln(w, `<p>dumped: <code id="`+clipID+`">stun-decode`, b64, `
				<button class="btn" data-clipboard-target="#`+clipID+`">copy</button>

`)
			fmt.Fprintln(w, `<p>crc64: <code>`, messageCRC64, `</code></p>`)
			fmt.Fprintln(w, `</div>`)
		}
	})

	var (
		addrSTUN = fmt.Sprintf(":%d", *portSTUN)
		addrHTTP = fmt.Sprintf("%s:%d", *hostHTTP, *portHTTP)
	)
	log.Println("Listening udp", addrSTUN)
	c, err := net.ListenPacket("udp", addrSTUN)
	if err != nil {
		log.Fatalf("Failed to bind udp on %s: %s", addrSTUN, err)
	}
	defer c.Close()

	// spawning storage garbage collector
	go messages.gc()

	// spawning STUN server
	go func(conn net.PacketConn) {
		log.Println("Started STUN server on", conn.LocalAddr())
		var (
			res = new(stun.Message)
			req = new(stun.Message)
			buf = make([]byte, 1024)
		)
		for {
			// ReadFrom c to buf
			n, addr, err := c.ReadFrom(buf)
			if err != nil {
				log.Fatalln("c.ReadFrom:", err)
			}
			log.Printf("udp: got packet len(%d) from %s", n, addr)
			// processing binding request
			if err = processUDPPacket(addr, buf[:n], req, res); err != nil {
				log.Println("failed to process UDP packet:", err, "from addr", addr)
			} else {
				log.Printf("stun: parsed message %q from %s", req, addr)
				if _, err = c.WriteTo(res.Raw, addr); err != nil {
					log.Println("failed to send packet:", err)
				}
			}
			res.Reset()
			req.Reset()
		}

	}(c)
	log.Println("Listening http", addrHTTP)
	log.Fatal(http.ListenAndServe(addrHTTP, nil))
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/nuveo/prest/cmd"
)

var Version = "0.1.2"

func main() {
	cmd.Execute()
}
<commit_msg>Fixed should have comment version var<commit_after>package main

import (
	"github.com/nuveo/prest/cmd"
)

// Version pREST
var Version = "0.1.3"

func main() {
	cmd.Execute()
}
<|endoftext|>


<commit_before>package main

import "io"
import "fmt"
import "net/http"
import "os"
import "os/exec"
import "flag"

import "github.com/gorilla/websocket"
import "github.com/kr/pty"

func start() (*exec.Cmd, *os.File) {
	var err error

	cmdString := "/bin/bash"
	cmd := exec.Command(cmdString)
	file, err := pty.Start(cmd)
	if err != nil {
		fmt.Println("Failed to start command: %s", err)
	}

	return cmd, file
}

func stop(pty *os.File, cmd *exec.Cmd) {
	pty.Close()
	cmd.Wait()
}

func ptyHandler(w http.ResponseWriter, r *http.Request) {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1,
		WriteBufferSize: 1,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println("Websocket upgrade failed: %s", err)
	}
	defer conn.Close()

	cmd, file := start()

	// Copy everything from the pty master to the websocket.
	go func() {
		buf := make([]byte, 256)
		// TODO: more graceful exit on socket close / process exit
		for {
			n, err := file.Read(buf)
			if err != nil {
				fmt.Println("Failed to read from pty master: %s", err)
				return
			}

			err = conn.WriteMessage(websocket.BinaryMessage, buf[0:n])

			if err != nil {
				fmt.Println("Failed to send %d bytes on websocket: %s", n, err)
				return
			}
		}
	}()

	// Read from the websocket, copying to the pty master.
	for {
		mt, payload, err := conn.ReadMessage()
		if err != nil {
			if err != io.EOF {
				fmt.Println("conn.ReadMessage failed: %s", err)
				return
			}
		}

		switch mt {
		case websocket.BinaryMessage:
			file.Write(payload)
		default:
			fmt.Println("Invalid message type %d", mt)
			return
		}
	}

	stop(file, cmd)
}

func main() {
	var addrFlag string
	flag.StringVar(&addrFlag, "addr", ":12061", "IP:PORT or :PORT address to listen on")

	flag.Parse()

	http.HandleFunc("/pty", ptyHandler)

	err := http.ListenAndServe(addrFlag, nil)
	if err != nil {
		fmt.Println("net.http could not listen on address '%s': %s", addrFlag, err)
	}
}
<commit_msg>Added pkg and sample code for winresize<commit_after>package main

import "io"
import "fmt"
import "net/http"
import "os"
import "os/exec"
import "flag"

import "github.com/gorilla/websocket"
import "github.com/kr/pty"
// import "github.com/creack/goterm/win"

func start() (*exec.Cmd, *os.File) {
	var err error

	cmdString := "/bin/bash"
	cmd := exec.Command(cmdString)
	f, err := pty.Start(cmd)
	if err != nil {
		fmt.Println("Failed to start command: %s", err)
	}

 //    if err := win.SetWinsize(f.Fd(), &win.Winsize{Height: 40, Width: 40}); err != nil {
 //        panic(err)
 //    }

 //    if size, err := win.GetWinsize(f.Fd()); err == nil {
 //        println(size.Height, size.Width)
 //    }

 //    if rows, cols, err := pty.Getsize(f); err == nil {
 //        println(rows, cols)
 //    }

	return cmd, f
}

func stop(pty *os.File, cmd *exec.Cmd) {
	pty.Close()
	cmd.Wait()
}

func ptyHandler(w http.ResponseWriter, r *http.Request) {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1,
		WriteBufferSize: 1,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println("Websocket upgrade failed: %s", err)
	}
	defer conn.Close()

	cmd, file := start()

	// Copy everything from the pty master to the websocket.
	go func() {
		buf := make([]byte, 256)
		// TODO: more graceful exit on socket close / process exit
		for {
			n, err := file.Read(buf)
			if err != nil {
				fmt.Println("Failed to read from pty master: %s", err)
				return
			}

			err = conn.WriteMessage(websocket.BinaryMessage, buf[0:n])

			if err != nil {
				fmt.Println("Failed to send %d bytes on websocket: %s", n, err)
				return
			}
		}
	}()

	// Read from the websocket, copying to the pty master.
	for {
		mt, payload, err := conn.ReadMessage()
		if err != nil {
			if err != io.EOF {
				fmt.Println("conn.ReadMessage failed: %s", err)
				return
			}
		}

		switch mt {
		case websocket.BinaryMessage:
			file.Write(payload)
		default:
			fmt.Println("Invalid message type %d", mt)
			return
		}
	}

	stop(file, cmd)
}

func main() {
	var addrFlag string
	flag.StringVar(&addrFlag, "addr", ":12061", "IP:PORT or :PORT address to listen on")

	flag.Parse()

	http.HandleFunc("/pty", ptyHandler)

	err := http.ListenAndServe(addrFlag, nil)
	if err != nil {
		fmt.Println("net.http could not listen on address '%s': %s", addrFlag, err)
	}
}
<|endoftext|>