<commit_before>package main

import (
	"bufio"
	"io"
	"net"
	"net/http"
	"os"
	"sync"

	"github.com/flynn/flynn/Godeps/_workspace/src/gopkg.in/inconshreveable/log15.v2"
	"github.com/flynn/flynn/discoverd/client"
	"github.com/flynn/flynn/logaggregator/snapshot"
	"github.com/flynn/flynn/logaggregator/utils"
	"github.com/flynn/flynn/pkg/syslog/rfc6587"
)

type Server struct {
	*Aggregator
	Cursors *HostCursors

	conf ServerConfig

	syslogListener net.Listener
	apiListener    net.Listener
	syslogWg       sync.WaitGroup

	hb discoverd.Heartbeater

	api      http.Handler
	shutdown chan struct{}
}

type ServerConfig struct {
	SyslogAddr, ApiAddr string

	ServiceName string
	Discoverd   *discoverd.Client
}

func NewServer(conf ServerConfig) *Server {
	a := NewAggregator()
	c := NewHostCursors()
	return &Server{
		Aggregator: a,
		Cursors:    c,
		conf:       conf,
		api:        apiHandler(a, c),
		shutdown:   make(chan struct{}),
	}
}

func (s *Server) Shutdown() {
	if s.hb != nil {
		// close discoverd service heartbeater
		if err := s.hb.Close(); err != nil {
			log15.Error("heartbeat shutdown error", "err", err)
		}
	}

	// shutdown listeners
	if s.syslogListener != nil {
		if err := s.syslogListener.Close(); err != nil {
			log15.Error("syslog listener shutdown error", "err", err)
		}
	}
	if s.apiListener != nil {
		if err := s.apiListener.Close(); err != nil {
			log15.Error("api listener shutdown error", "err", err)
		}
	}

	// close syslog client connections
	close(s.shutdown)
	s.syslogWg.Wait()

	// shutdown aggregator
	s.Aggregator.Shutdown()
}

func (s *Server) LoadSnapshotFile(path string) error {
	f, err := os.Open(path)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	defer f.Close()
	return s.LoadSnapshot(f)
}

func (s *Server) LoadSnapshot(r io.Reader) error {
	sc := snapshot.NewScanner(r)
	for sc.Scan() {
		cursor, err := utils.ParseHostCursor(sc.Message)
		if err != nil {
			return err
		}
		s.Cursors.Update(string(sc.Message.Hostname), cursor)
		s.Aggregator.Feed(sc.Message)
	}
	return sc.Err()
}

func (s *Server) WriteSnapshotFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	return snapshot.WriteTo(s.Aggregator.ReadAll(), f)
}

func (s *Server) SyslogAddr() net.Addr {
	return s.syslogListener.Addr()
}

func (s *Server) Start() error {
	var err error
	s.syslogListener, err = net.Listen("tcp", s.conf.SyslogAddr)
	if err != nil {
		return err
	}

	s.apiListener, err = net.Listen("tcp", s.conf.ApiAddr)
	if err != nil {
		return err
	}

	if s.conf.Discoverd != nil {
		s.hb, err = s.conf.Discoverd.AddServiceAndRegister(s.conf.ServiceName, s.conf.SyslogAddr)
		if err != nil {
			return err
		}
	}

	go s.runSyslog()
	go http.Serve(s.apiListener, s.api)

	return nil
}

func (s *Server) runSyslog() {
	for {
		conn, err := s.syslogListener.Accept()
		if err != nil {
			return
		}

		s.syslogWg.Add(1)
		go func(c net.Conn) {
			defer s.syslogWg.Done()
			s.drainSyslogConn(c)
		}(conn)
	}
}

func (s *Server) drainSyslogConn(conn net.Conn) {
	connDone := make(chan struct{})
	defer close(connDone)

	go func() {
		select {
		case <-connDone:
		case <-s.shutdown:
		}
		conn.Close()
	}()

	sc := bufio.NewScanner(conn)
	sc.Split(rfc6587.Split)
	for sc.Scan() {
		msgBytes := sc.Bytes()
		// slice in msgBytes could get modified on next Scan(), need to copy it
		msgCopy := make([]byte, len(msgBytes))
		copy(msgCopy, msgBytes)

		msg, cursor, err := utils.ParseMessage(msgCopy)
		if err != nil {
			log15.Error("rfc5424 parse error", "err", err)
		} else {
			s.Cursors.Update(string(msg.Hostname), cursor)
			s.Aggregator.Feed(msg)
		}
	}
}

func NewHostCursors() *HostCursors {
	return &HostCursors{Data: make(map[string]*utils.HostCursor)}
}

// HostCursors are used to keep track of what messages a host has sent when
// resuming log streaming to a new server or after an interruption .
type HostCursors struct {
	mtx  sync.Mutex
	Data map[string]*utils.HostCursor
}

func (h *HostCursors) Update(id string, other *utils.HostCursor) {
	h.mtx.Lock()
	defer h.mtx.Unlock()
	curr, ok := h.Data[id]
	if !ok || other.After(*curr) {
		h.Data[id] = other
	}
}

func (h *HostCursors) Get() map[string]*utils.HostCursor {
	h.mtx.Lock()
	defer h.mtx.Unlock()
	res := make(map[string]*utils.HostCursor, len(h.Data))
	for k, v := range h.Data {
		res[k] = v
	}
	return res
}
<commit_msg>logaggregator: Fix WaitGroup race during shutdown<commit_after>package main

import (
	"bufio"
	"io"
	"net"
	"net/http"
	"os"
	"sync"

	"github.com/flynn/flynn/Godeps/_workspace/src/gopkg.in/inconshreveable/log15.v2"
	"github.com/flynn/flynn/discoverd/client"
	"github.com/flynn/flynn/logaggregator/snapshot"
	"github.com/flynn/flynn/logaggregator/utils"
	"github.com/flynn/flynn/pkg/syslog/rfc6587"
)

type Server struct {
	*Aggregator
	Cursors *HostCursors

	conf ServerConfig

	syslogListener net.Listener
	apiListener    net.Listener
	syslogWg       sync.WaitGroup
	syslogDone     chan struct{}

	hb discoverd.Heartbeater

	api      http.Handler
	shutdown chan struct{}
}

type ServerConfig struct {
	SyslogAddr, ApiAddr string

	ServiceName string
	Discoverd   *discoverd.Client
}

func NewServer(conf ServerConfig) *Server {
	a := NewAggregator()
	c := NewHostCursors()
	return &Server{
		Aggregator: a,
		Cursors:    c,
		conf:       conf,
		api:        apiHandler(a, c),
		shutdown:   make(chan struct{}),
		syslogDone: make(chan struct{}),
	}
}

func (s *Server) Shutdown() {
	if s.hb != nil {
		// close discoverd service heartbeater
		if err := s.hb.Close(); err != nil {
			log15.Error("heartbeat shutdown error", "err", err)
		}
	}

	// shutdown listeners
	if s.syslogListener != nil {
		if err := s.syslogListener.Close(); err != nil {
			log15.Error("syslog listener shutdown error", "err", err)
		}
		<-s.syslogDone
	}
	if s.apiListener != nil {
		if err := s.apiListener.Close(); err != nil {
			log15.Error("api listener shutdown error", "err", err)
		}
	}

	// close syslog client connections
	close(s.shutdown)
	s.syslogWg.Wait()

	// shutdown aggregator
	s.Aggregator.Shutdown()
}

func (s *Server) LoadSnapshotFile(path string) error {
	f, err := os.Open(path)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	defer f.Close()
	return s.LoadSnapshot(f)
}

func (s *Server) LoadSnapshot(r io.Reader) error {
	sc := snapshot.NewScanner(r)
	for sc.Scan() {
		cursor, err := utils.ParseHostCursor(sc.Message)
		if err != nil {
			return err
		}
		s.Cursors.Update(string(sc.Message.Hostname), cursor)
		s.Aggregator.Feed(sc.Message)
	}
	return sc.Err()
}

func (s *Server) WriteSnapshotFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	return snapshot.WriteTo(s.Aggregator.ReadAll(), f)
}

func (s *Server) SyslogAddr() net.Addr {
	return s.syslogListener.Addr()
}

func (s *Server) Start() error {
	var err error
	s.syslogListener, err = net.Listen("tcp", s.conf.SyslogAddr)
	if err != nil {
		return err
	}

	s.apiListener, err = net.Listen("tcp", s.conf.ApiAddr)
	if err != nil {
		return err
	}

	if s.conf.Discoverd != nil {
		s.hb, err = s.conf.Discoverd.AddServiceAndRegister(s.conf.ServiceName, s.conf.SyslogAddr)
		if err != nil {
			return err
		}
	}

	go s.runSyslog()
	go http.Serve(s.apiListener, s.api)

	return nil
}

func (s *Server) runSyslog() {
	defer close(s.syslogDone)
	for {
		conn, err := s.syslogListener.Accept()
		if err != nil {
			return
		}

		s.syslogWg.Add(1)
		go func(c net.Conn) {
			defer s.syslogWg.Done()
			s.drainSyslogConn(c)
		}(conn)
	}
}

func (s *Server) drainSyslogConn(conn net.Conn) {
	connDone := make(chan struct{})
	defer close(connDone)

	go func() {
		select {
		case <-connDone:
		case <-s.shutdown:
		}
		conn.Close()
	}()

	sc := bufio.NewScanner(conn)
	sc.Split(rfc6587.Split)
	for sc.Scan() {
		msgBytes := sc.Bytes()
		// slice in msgBytes could get modified on next Scan(), need to copy it
		msgCopy := make([]byte, len(msgBytes))
		copy(msgCopy, msgBytes)

		msg, cursor, err := utils.ParseMessage(msgCopy)
		if err != nil {
			log15.Error("rfc5424 parse error", "err", err)
		} else {
			s.Cursors.Update(string(msg.Hostname), cursor)
			s.Aggregator.Feed(msg)
		}
	}
}

func NewHostCursors() *HostCursors {
	return &HostCursors{Data: make(map[string]*utils.HostCursor)}
}

// HostCursors are used to keep track of what messages a host has sent when
// resuming log streaming to a new server or after an interruption .
type HostCursors struct {
	mtx  sync.Mutex
	Data map[string]*utils.HostCursor
}

func (h *HostCursors) Update(id string, other *utils.HostCursor) {
	h.mtx.Lock()
	defer h.mtx.Unlock()
	curr, ok := h.Data[id]
	if !ok || other.After(*curr) {
		h.Data[id] = other
	}
}

func (h *HostCursors) Get() map[string]*utils.HostCursor {
	h.mtx.Lock()
	defer h.mtx.Unlock()
	res := make(map[string]*utils.HostCursor, len(h.Data))
	for k, v := range h.Data {
		res[k] = v
	}
	return res
}
<|endoftext|>


<commit_before>package logger

import (
	"bytes"
	"fmt"
	ct "github.com/daviddengcn/go-colortext"
	"github.com/gosuri/uilive"
	"github.com/op/go-logging"
	"time"
)

type coloredLogger struct {
	writer *uilive.Writer

	sysoutBuffer *bytes.Buffer
}

func newColoredConsoleWriter() *coloredLogger {
	return &coloredLogger{writer: uilive.New(), sysoutBuffer: new(bytes.Buffer)}
}

func (cLogger *coloredLogger) writeSysoutBuffer(text string) {
	cLogger.sysoutBuffer.WriteString(text)
}

func (cLogger *coloredLogger) Info(msg string, args ...interface{}) {
	Log.Info(msg, args...)
	cLogger.ConsoleWrite(msg, args...)
}

func (cLogger *coloredLogger) Debug(msg string, args ...interface{}) {
	Log.Info(msg, args...)
	if level == logging.DEBUG {
		cLogger.ConsoleWrite(msg, args...)
	}
}

func (cLogger *coloredLogger) ConsoleWrite(msg string, args ...interface{}) {
	fmt.Println(fmt.Sprintf(msg, args...))
}

func (cLogger *coloredLogger) Spec(heading string) {
	msg := formatSpec(heading)
	cLogger.Info(msg)
}

func (cLogger *coloredLogger) ScenarioStart(scenarioHeading string) {
	msg := formatScenario(scenarioHeading)
	Log.Info(msg)
	if level == logging.INFO {
		cLogger.start(indent(msg, scenarioIndentation))
	} else {
		fmt.Println(indent(msg, scenarioIndentation))
		fmt.Println()
	}
}

func (cLogger *coloredLogger) ScenarioEnd(scenarioHeading string, failed bool) {
	msg := formatScenario(scenarioHeading)
	if level == logging.INFO {
		cLogger.end(indent(msg, scenarioIndentation), failed)
	}
}

func (cLogger *coloredLogger) StepStart(stepText string) {
	Log.Debug(stepText)
	if level == logging.DEBUG {
		cLogger.start(indent(stepText, stepIndentation))
	}
}

func (cLogger *coloredLogger) StepEnd(stepText string, failed bool) {
	if level == logging.DEBUG {
		cLogger.end(indent(stepText, stepIndentation), failed)
	}
}

func (cLogger *coloredLogger) Table(table string) {
}

func (cLogger *coloredLogger) start(text string) {
	cLogger.writer.Start()
	ct.Foreground(ct.Cyan, true)
	fmt.Fprintln(cLogger.writer, text)
	cLogger.writer.Flush()
	ct.ResetColor()
}

func (cLogger *coloredLogger) end(text string, failed bool) {
	time.Sleep(time.Millisecond * 200)
	if failed {
		ct.Foreground(ct.Red, true)
	} else {
		ct.Foreground(ct.Green, true)
	}
	cLogger.writer.Flush()
	fmt.Fprintln(cLogger.writer, text)

	cLogger.writer.Stop()
	ct.ResetColor()
	fmt.Println()
}
<commit_msg>starting the console writer listener at scenario level<commit_after>package logger

import (
	"bytes"
	"fmt"
	ct "github.com/daviddengcn/go-colortext"
	"github.com/gosuri/uilive"
	"github.com/op/go-logging"
	"strings"
)

type coloredLogger struct {
	writer *uilive.Writer

	sysoutBuffer *bytes.Buffer
}

func newColoredConsoleWriter() *coloredLogger {
	return &coloredLogger{writer: uilive.New(), sysoutBuffer: new(bytes.Buffer)}
}

func (cLogger *coloredLogger) writeSysoutBuffer(text string) {
	cLogger.sysoutBuffer.WriteString(text)
}

func (cLogger *coloredLogger) Info(msg string, args ...interface{}) {
	Log.Info(msg, args...)
	cLogger.ConsoleWrite(msg, args...)
}

func (cLogger *coloredLogger) Debug(msg string, args ...interface{}) {
	Log.Info(msg, args...)
	if level == logging.DEBUG {
		cLogger.ConsoleWrite(msg, args...)
	}
}

func (cLogger *coloredLogger) ConsoleWrite(msg string, args ...interface{}) {
	fmt.Println(fmt.Sprintf(msg, args...))
}

func (cLogger *coloredLogger) Spec(heading string) {
	msg := formatSpec(heading)
	cLogger.Info(msg)
}

func (cLogger *coloredLogger) ScenarioStart(scenarioHeading string) {
	msg := formatScenario(scenarioHeading)
	Log.Info(msg)
	cLogger.writer.Start()
	if level == logging.INFO {
		cLogger.start(indent(msg, scenarioIndentation))
	} else {

		fmt.Println(indent(msg, scenarioIndentation))
		fmt.Println()
	}
}

func (cLogger *coloredLogger) ScenarioEnd(scenarioHeading string, failed bool) {
	msg := formatScenario(scenarioHeading)
	if level == logging.INFO {
		cLogger.end(indent(msg, scenarioIndentation), failed)
	}
	cLogger.writer.Stop()
}

func (cLogger *coloredLogger) StepStart(stepText string) {
	Log.Debug(stepText)
	if level == logging.DEBUG {
		cLogger.start(indent(stepText, stepIndentation))
	}
}

func (cLogger *coloredLogger) StepEnd(stepText string, failed bool) {
	if level == logging.DEBUG {
		cLogger.end(indent(stepText, stepIndentation), failed)
	}
}

func (cLogger *coloredLogger) Table(table string) {
}

func (cLogger *coloredLogger) start(text string) {
	ct.Foreground(ct.Cyan, true)
	fragments := strings.Split(text, "\n")
	for i := 0; i < len(fragments); i++ {
		fmt.Fprintln(cLogger.writer, fragments[i])
	}
	cLogger.writer.Flush()
	ct.ResetColor()
}

func (cLogger *coloredLogger) end(text string, failed bool) {
	if failed {
		ct.Foreground(ct.Red, true)
	} else {
		ct.Foreground(ct.Green, true)
	}
	fragments := strings.Split(text, "\n")
	for i := 0; i < len(fragments); i++ {
		fmt.Fprintln(cLogger.writer, fragments[i])
	}
	cLogger.writer.Flush()
	ct.ResetColor()
	fmt.Println()
}
<|endoftext|>


<commit_before>package login

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/Bplotka/go-httpt/rt"
	"github.com/Bplotka/oidc"
	"github.com/Bplotka/oidc/testing"
	"github.com/stretchr/testify/suite"
	"github.com/stretchr/testify/mock"
)

const (
	testBindAddress  = "http://127.0.0.1:0/something/callback"
	testSubject      = "subject1"
	testClientID     = "clientID1"
	testClientSecret = "secret1"
	testNonce        = "nonce1"
)

var (
	testToken = oidc.Token{
		AccessToken:  "access1",
		RefreshToken: "refresh1",
		IDToken:      "idtoken1",
	}
)

type TokenSourceTestSuite struct {
	suite.Suite

	testCfg     Config
	testOIDCCfg OIDCConfig

	cache      *MockCache
	oidcSource *OIDCTokenSource

	provider *oidc_testing.Provider

	closeSrv func()
}

func (s *TokenSourceTestSuite) SetupSuite() {
	s.provider = &oidc_testing.Provider{}
	s.provider.Setup(s.T())
	s.provider.MockDiscoveryCall()

	s.testOIDCCfg = OIDCConfig{
		Provider: s.provider.IssuerURL,

		ClientID:     testClientID,
		ClientSecret: testClientSecret,
		Scopes:       []string{oidc.ScopeOpenID, oidc.ScopeEmail},
	}
	s.testCfg = Config{
		NonceCheck: true,
	}

	s.cache = new(MockCache)

	callbackSrv, closeSrv, err := NewServer(testBindAddress)
	s.Require().NoError(err)

	s.closeSrv = closeSrv

	oidcClient, err := oidc.NewClient(s.provider.Context(), s.testOIDCCfg.Provider)
	s.Require().NoError(err)

	s.oidcSource = &OIDCTokenSource{
		ctx:    s.provider.Context(),
		logger: log.New(os.Stdout, "", 0),
		cfg:    s.testCfg,

		oidcClient:  oidcClient,
		cache:       s.cache,
		openBrowser: openBrowser,
		callbackSrv: callbackSrv,
		nonce:       testNonce,
	}
}

func (s *TokenSourceTestSuite) TearDownSuite() {
	s.closeSrv()
}

func (s *TokenSourceTestSuite) SetupTest() {
	s.provider.Mock().Reset()

	s.oidcSource.openBrowser = func(string) error {
		s.T().Errorf("OpenBrowser Not mocked")
		s.T().FailNow()
		return nil
	}
	s.oidcSource.genRandToken = func() string {
		s.T().Errorf("GenState Not mocked")
		s.T().FailNow()
		return ""
	}

	s.cache = new(MockCache)
	s.cache.On("Config").Return(s.testOIDCCfg)
	s.oidcSource.cache = s.cache
}

func TestTokenSourceTestSuite(t *testing.T) {
	suite.Run(t, &TokenSourceTestSuite{})
}

func (s *TokenSourceTestSuite) Test_CacheOK() {
	idToken, jwkSetJSON := s.provider.NewIDToken(testClientID, testSubject, s.oidcSource.nonce)
	expectedToken := testToken
	expectedToken.IDToken = idToken
	s.cache.On("Token").Return(&expectedToken, nil)

	s.provider.MockPubKeysCall(jwkSetJSON)

	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(expectedToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

// stripArgFromURL strips out arg value from URL.
func stripArgFromURL(arg string, urlToStrip string) (string, error) {
	var argValue string
	splittedURL := strings.Split(urlToStrip, "&")
	for _, a := range splittedURL {
		if !strings.HasPrefix(a, arg+"=") {
			continue
		}
		splittedArg := strings.Split(a, "=")
		if len(splittedArg) != 2 {
			return "", errors.New("More or less than two args after splitting by `=`")
		}
		var err error
		argValue, err = url.QueryUnescape(splittedArg[1])
		if err != nil {
			return "", err
		}
	}
	if argValue == "" {
		return "", fmt.Errorf("%s not found in given URL.", arg)
	}
	return argValue, nil
}

func (s *TokenSourceTestSuite) callSuccessfulCallback(expectedWord string) func(string) error {
	return func(urlToGet string) error {
		redirectURL, err := stripArgFromURL("redirect_uri", urlToGet)
		s.Require().NoError(err)

		s.Equal(fmt.Sprintf(
			"https://issuer.org/auth1?client_id=%s&nonce=%s&redirect_uri=%s&response_type=code&scope=%s&state=%s",
			testClientID,
			expectedWord,
			url.QueryEscape(redirectURL),
			strings.Join(s.testOIDCCfg.Scopes, "+"),
			expectedWord,
		), urlToGet)

		t := oidc.TokenResponse{
			AccessToken:  testToken.AccessToken,
			RefreshToken: testToken.RefreshToken,
			IDToken:      testToken.IDToken,
			TokenType:    "Bearer",
		}
		tokenJSON, err := json.Marshal(t)
		s.Require().NoError(err)

		s.provider.Mock().Push(rt.JSONResponseFunc(http.StatusOK, tokenJSON))

		go func() {
			// Perform actual request in go routine.
			req, err := http.NewRequest("GET", fmt.Sprintf(
				"%s?code=%s&state=%s",
				redirectURL,
				"code1",
				expectedWord,
			), nil)
			s.Require().NoError(err)

			u, err := url.Parse(redirectURL)
			s.Require().NoError(err)
			for i := 0; i <= 5; i++ {
				_, err = net.Dial("tcp", u.Host)
				if err == nil {
					break
				}
				time.Sleep(100 * time.Millisecond)
			}
			s.Require().NoError(err, "Server should be able to start and listen on provided address.")

			res, err := http.DefaultClient.Do(req)
			s.Require().NoError(err)

			s.Equal(http.StatusOK, res.StatusCode)
		}()
		return nil
	}
}

func (s *TokenSourceTestSuite) Test_CacheErr_NewToken_OKCallback() {
	s.cache.On("Token").Return(nil, errors.New("test_err"))
	s.cache.On("SaveToken", &testToken).Return(nil)

	const expectedWord = "secret_token"
	s.oidcSource.genRandToken = func() string {
		return expectedWord
	}

	s.oidcSource.openBrowser = s.callSuccessfulCallback(expectedWord)
	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(testToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_CacheEmpty_NewToken_OKCallback() {
	s.cache.On("Token").Return(nil, nil)
	s.cache.On("SaveToken", &testToken).Return(nil)

	const expectedWord = "secret_token"
	s.oidcSource.genRandToken = func() string {
		return expectedWord
	}

	s.oidcSource.openBrowser = s.callSuccessfulCallback(expectedWord)
	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(testToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_IDTokenWrongNonce_RefreshToken_OK() {
	idToken, jwkSetJSON := s.provider.NewIDToken(testClientID, testSubject, "wrongNonce")
	invalidToken := testToken
	invalidToken.IDToken = idToken
	s.cache.On("Token").Return(&invalidToken, nil)

	idTokenOkNonce, jwkSetJSON2 := s.provider.NewIDToken(testClientID, testSubject, s.oidcSource.nonce)
	expectedToken := invalidToken
	expectedToken.IDToken = idTokenOkNonce
	s.cache.On("SaveToken", &expectedToken).Return(nil)

	// For first verification inside OIDC TokenSource.
	s.provider.MockPubKeysCall(jwkSetJSON)

	// OK Refresh response.
	t := oidc.TokenResponse{
		AccessToken:  expectedToken.AccessToken,
		RefreshToken: expectedToken.RefreshToken,
		IDToken:      expectedToken.IDToken,
		TokenType:    "Bearer",
	}
	tokenJSON, err := json.Marshal(t)
	s.Require().NoError(err)

	s.provider.Mock().Push(rt.JSONResponseFunc(http.StatusOK, tokenJSON))

	// For 2th verification inside reuse TokenSource.
	s.provider.MockPubKeysCall(jwkSetJSON2)

	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(expectedToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_IDTokenWrongNonce_RefreshTokenErr_NewToken_OK() {
	idToken, jwkSetJSON := s.provider.NewIDToken(testClientID, testSubject, "wrongNonce")
	invalidToken := testToken
	invalidToken.IDToken = idToken
	s.cache.On("Token").Return(&invalidToken, nil)
	s.cache.On("SaveToken", &testToken).Return(nil)

	// For first verification inside OIDC TokenSource.
	s.provider.MockPubKeysCall(jwkSetJSON)

	s.provider.Mock().Push(rt.JSONResponseFunc(http.StatusBadRequest, []byte(`{"error": "bad_request"}`)))

	const expectedWord = "secret_token"
	s.oidcSource.genRandToken = func() string {
		return expectedWord
	}
	s.oidcSource.openBrowser = s.callSuccessfulCallback(expectedWord)

	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(testToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_CacheEmpty_NewToken_ErrCallback() {
	s.cache.On("Token").Return(nil, nil)

	const expectedWord = "secret_token"
	s.oidcSource.genRandToken = func() string {
		return expectedWord
	}

	s.oidcSource.openBrowser = func(urlToGet string) error {
		redirectURL, err := stripArgFromURL("redirect_uri", urlToGet)
		s.Require().NoError(err)

		s.Equal(fmt.Sprintf(
			"https://issuer.org/auth1?client_id=%s&nonce=%s&redirect_uri=%s&response_type=code&scope=%s&state=%s",
			testClientID,
			expectedWord,
			url.QueryEscape(redirectURL),
			strings.Join(s.testOIDCCfg.Scopes, "+"),
			expectedWord,
		), urlToGet)

		s.provider.Mock().Push(rt.JSONResponseFunc(http.StatusGatewayTimeout, []byte(`{"error": "temporary unavailable"}`)))

		go func() {
			req, err := http.NewRequest("GET", fmt.Sprintf(
				"%s?code=%s&state=%s",
				redirectURL,
				"code1",
				expectedWord,
			), nil)
			s.Require().NoError(err)

			u, err := url.Parse(redirectURL)
			s.Require().NoError(err)
			for i := 0; i <= 5; i++ {
				_, err = net.Dial("tcp", u.Host)
				if err == nil {
					break
				}
				time.Sleep(100 * time.Millisecond)
			}
			s.Require().NoError(err)

			res, err := http.DefaultClient.Do(req)
			s.Require().NoError(err)

			// Still it should be ok.
			s.Equal(http.StatusOK, res.StatusCode)
		}()

		return nil
	}

	_, err := s.oidcSource.OIDCToken()
	s.Require().Error(err)
	s.Equal("Failed to obtain new token. Err: oidc: Callback error: oauth2: cannot fetch token: \nResponse: {\"error\": \"temporary unavailable\"}", err.Error())

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_ClearIDToken_ClearOnlyIDToken() {
	token := oidc.Token{
		AccessToken: "accessToken",
		IDToken: "idToken",
		RefreshToken: "refreshToken",
	}
	// Just to satisfy mock.
	s.cache.Config()

	s.cache.On("Token").Return(&token, nil)
	s.cache.On("SaveToken", mock.Anything).Run(func(a mock.Arguments){
		t, ok := a.Get(0).(*oidc.Token)
		s.Require().True(ok)

		s.Assert().Equal(token.AccessToken, t.AccessToken)
		s.Assert().Empty(t.IDToken)
		s.Assert().Equal(token.RefreshToken, t.RefreshToken)
	}).Return(nil)

	resetDone := true
	clear := s.oidcSource.clearIDToken(func(){resetDone = true})
	s.Require().NoError(clear())

	s.cache.AssertExpectations(s.T())
}<commit_msg>Fixed test.<commit_after>package login

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/Bplotka/go-httpt/rt"
	"github.com/Bplotka/oidc"
	"github.com/Bplotka/oidc/testing"
	"github.com/stretchr/testify/suite"
	"github.com/stretchr/testify/mock"
)

const (
	testBindAddress  = "http://127.0.0.1:0/something/callback"
	testSubject      = "subject1"
	testClientID     = "clientID1"
	testClientSecret = "secret1"
	testNonce        = "nonce1"
)

var (
	testToken = oidc.Token{
		AccessToken:  "access1",
		RefreshToken: "refresh1",
		IDToken:      "idtoken1",
	}
)

type TokenSourceTestSuite struct {
	suite.Suite

	testCfg     Config
	testOIDCCfg OIDCConfig

	cache      *MockCache
	oidcSource *OIDCTokenSource

	provider *oidc_testing.Provider

	closeSrv func()
}

func (s *TokenSourceTestSuite) SetupSuite() {
	s.provider = &oidc_testing.Provider{}
	s.provider.Setup(s.T())
	s.provider.MockDiscoveryCall()

	s.testOIDCCfg = OIDCConfig{
		Provider: s.provider.IssuerURL,

		ClientID:     testClientID,
		ClientSecret: testClientSecret,
		Scopes:       []string{oidc.ScopeOpenID, oidc.ScopeEmail},
	}
	s.testCfg = Config{
		NonceCheck: true,
	}

	s.cache = new(MockCache)

	callbackSrv, closeSrv, err := NewServer(testBindAddress)
	s.Require().NoError(err)

	s.closeSrv = closeSrv

	oldKeySetExpiration := oidc.DefaultKeySetExpiration
	oidc.DefaultKeySetExpiration = 0 * time.Second
	defer func() {
		oidc.DefaultKeySetExpiration = oldKeySetExpiration
	}()
	oidcClient, err := oidc.NewClient(s.provider.Context(), s.testOIDCCfg.Provider)
	s.Require().NoError(err)

	s.oidcSource = &OIDCTokenSource{
		ctx:    s.provider.Context(),
		logger: log.New(os.Stdout, "", 0),
		cfg:    s.testCfg,

		oidcClient:  oidcClient,
		cache:       s.cache,
		openBrowser: openBrowser,
		callbackSrv: callbackSrv,
		nonce:       testNonce,
	}
}

func (s *TokenSourceTestSuite) TearDownSuite() {
	s.closeSrv()
}

func (s *TokenSourceTestSuite) SetupTest() {
	s.provider.Mock().Reset()

	s.oidcSource.openBrowser = func(string) error {
		s.T().Errorf("OpenBrowser Not mocked")
		s.T().FailNow()
		return nil
	}
	s.oidcSource.genRandToken = func() string {
		s.T().Errorf("GenState Not mocked")
		s.T().FailNow()
		return ""
	}

	s.cache = new(MockCache)
	s.cache.On("Config").Return(s.testOIDCCfg)
	s.oidcSource.cache = s.cache
}

func TestTokenSourceTestSuite(t *testing.T) {
	suite.Run(t, &TokenSourceTestSuite{})
}

func (s *TokenSourceTestSuite) Test_CacheOK() {
	idToken, jwkSetJSON := s.provider.NewIDToken(testClientID, testSubject, s.oidcSource.nonce)
	expectedToken := testToken
	expectedToken.IDToken = idToken
	s.cache.On("Token").Return(&expectedToken, nil)

	s.provider.MockPubKeysCall(jwkSetJSON)

	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(expectedToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

// stripArgFromURL strips out arg value from URL.
func stripArgFromURL(arg string, urlToStrip string) (string, error) {
	var argValue string
	splittedURL := strings.Split(urlToStrip, "&")
	for _, a := range splittedURL {
		if !strings.HasPrefix(a, arg+"=") {
			continue
		}
		splittedArg := strings.Split(a, "=")
		if len(splittedArg) != 2 {
			return "", errors.New("More or less than two args after splitting by `=`")
		}
		var err error
		argValue, err = url.QueryUnescape(splittedArg[1])
		if err != nil {
			return "", err
		}
	}
	if argValue == "" {
		return "", fmt.Errorf("%s not found in given URL.", arg)
	}
	return argValue, nil
}

func (s *TokenSourceTestSuite) callSuccessfulCallback(expectedWord string) func(string) error {
	return func(urlToGet string) error {
		redirectURL, err := stripArgFromURL("redirect_uri", urlToGet)
		s.Require().NoError(err)

		s.Equal(fmt.Sprintf(
			"https://issuer.org/auth1?client_id=%s&nonce=%s&redirect_uri=%s&response_type=code&scope=%s&state=%s",
			testClientID,
			expectedWord,
			url.QueryEscape(redirectURL),
			strings.Join(s.testOIDCCfg.Scopes, "+"),
			expectedWord,
		), urlToGet)

		t := oidc.TokenResponse{
			AccessToken:  testToken.AccessToken,
			RefreshToken: testToken.RefreshToken,
			IDToken:      testToken.IDToken,
			TokenType:    "Bearer",
		}
		tokenJSON, err := json.Marshal(t)
		s.Require().NoError(err)

		s.provider.Mock().Push(rt.JSONResponseFunc(http.StatusOK, tokenJSON))

		go func() {
			// Perform actual request in go routine.
			req, err := http.NewRequest("GET", fmt.Sprintf(
				"%s?code=%s&state=%s",
				redirectURL,
				"code1",
				expectedWord,
			), nil)
			s.Require().NoError(err)

			u, err := url.Parse(redirectURL)
			s.Require().NoError(err)
			for i := 0; i <= 5; i++ {
				_, err = net.Dial("tcp", u.Host)
				if err == nil {
					break
				}
				time.Sleep(100 * time.Millisecond)
			}
			s.Require().NoError(err, "Server should be able to start and listen on provided address.")

			res, err := http.DefaultClient.Do(req)
			s.Require().NoError(err)

			s.Equal(http.StatusOK, res.StatusCode)
		}()
		return nil
	}
}

func (s *TokenSourceTestSuite) Test_CacheErr_NewToken_OKCallback() {
	s.cache.On("Token").Return(nil, errors.New("test_err"))
	s.cache.On("SaveToken", &testToken).Return(nil)

	const expectedWord = "secret_token"
	s.oidcSource.genRandToken = func() string {
		return expectedWord
	}

	s.oidcSource.openBrowser = s.callSuccessfulCallback(expectedWord)
	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(testToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_CacheEmpty_NewToken_OKCallback() {
	s.cache.On("Token").Return(nil, nil)
	s.cache.On("SaveToken", &testToken).Return(nil)

	const expectedWord = "secret_token"
	s.oidcSource.genRandToken = func() string {
		return expectedWord
	}

	s.oidcSource.openBrowser = s.callSuccessfulCallback(expectedWord)
	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(testToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_IDTokenWrongNonce_RefreshToken_OK() {
	idToken, jwkSetJSON := s.provider.NewIDToken(testClientID, testSubject, "wrongNonce")
	invalidToken := testToken
	invalidToken.IDToken = idToken
	s.cache.On("Token").Return(&invalidToken, nil)

	idTokenOkNonce, jwkSetJSON2 := s.provider.NewIDToken(testClientID, testSubject, s.oidcSource.nonce)
	expectedToken := invalidToken
	expectedToken.IDToken = idTokenOkNonce
	s.cache.On("SaveToken", &expectedToken).Return(nil)

	// For first verification inside OIDC TokenSource.
	s.provider.MockPubKeysCall(jwkSetJSON)

	// OK Refresh response.
	t := oidc.TokenResponse{
		AccessToken:  expectedToken.AccessToken,
		RefreshToken: expectedToken.RefreshToken,
		IDToken:      expectedToken.IDToken,
		TokenType:    "Bearer",
	}
	tokenJSON, err := json.Marshal(t)
	s.Require().NoError(err)

	s.provider.Mock().Push(rt.JSONResponseFunc(http.StatusOK, tokenJSON))

	// For 2th verification inside reuse TokenSource.
	s.provider.MockPubKeysCall(jwkSetJSON2)

	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(expectedToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_IDTokenWrongNonce_RefreshTokenErr_NewToken_OK() {
	idToken, jwkSetJSON := s.provider.NewIDToken(testClientID, testSubject, "wrongNonce")
	invalidToken := testToken
	invalidToken.IDToken = idToken
	s.cache.On("Token").Return(&invalidToken, nil)
	s.cache.On("SaveToken", &testToken).Return(nil)

	// For first verification inside OIDC TokenSource.
	s.provider.MockPubKeysCall(jwkSetJSON)

	s.provider.Mock().Push(rt.JSONResponseFunc(http.StatusBadRequest, []byte(`{"error": "bad_request"}`)))

	const expectedWord = "secret_token"
	s.oidcSource.genRandToken = func() string {
		return expectedWord
	}
	s.oidcSource.openBrowser = s.callSuccessfulCallback(expectedWord)

	token, err := s.oidcSource.OIDCToken()
	s.Require().NoError(err)

	s.Equal(testToken, *token)

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_CacheEmpty_NewToken_ErrCallback() {
	s.cache.On("Token").Return(nil, nil)

	const expectedWord = "secret_token"
	s.oidcSource.genRandToken = func() string {
		return expectedWord
	}

	s.oidcSource.openBrowser = func(urlToGet string) error {
		redirectURL, err := stripArgFromURL("redirect_uri", urlToGet)
		s.Require().NoError(err)

		s.Equal(fmt.Sprintf(
			"https://issuer.org/auth1?client_id=%s&nonce=%s&redirect_uri=%s&response_type=code&scope=%s&state=%s",
			testClientID,
			expectedWord,
			url.QueryEscape(redirectURL),
			strings.Join(s.testOIDCCfg.Scopes, "+"),
			expectedWord,
		), urlToGet)

		s.provider.Mock().Push(rt.JSONResponseFunc(http.StatusGatewayTimeout, []byte(`{"error": "temporary unavailable"}`)))

		go func() {
			req, err := http.NewRequest("GET", fmt.Sprintf(
				"%s?code=%s&state=%s",
				redirectURL,
				"code1",
				expectedWord,
			), nil)
			s.Require().NoError(err)

			u, err := url.Parse(redirectURL)
			s.Require().NoError(err)
			for i := 0; i <= 5; i++ {
				_, err = net.Dial("tcp", u.Host)
				if err == nil {
					break
				}
				time.Sleep(100 * time.Millisecond)
			}
			s.Require().NoError(err)

			res, err := http.DefaultClient.Do(req)
			s.Require().NoError(err)

			// Still it should be ok.
			s.Equal(http.StatusOK, res.StatusCode)
		}()

		return nil
	}

	_, err := s.oidcSource.OIDCToken()
	s.Require().Error(err)
	s.Equal("Failed to obtain new token. Err: oidc: Callback error: oauth2: cannot fetch token: \nResponse: {\"error\": \"temporary unavailable\"}", err.Error())

	s.cache.AssertExpectations(s.T())
	s.Equal(0, s.provider.Mock().Len())
}

func (s *TokenSourceTestSuite) Test_ClearIDToken_ClearOnlyIDToken() {
	token := oidc.Token{
		AccessToken: "accessToken",
		IDToken: "idToken",
		RefreshToken: "refreshToken",
	}
	// Just to satisfy mock.
	s.cache.Config()

	s.cache.On("Token").Return(&token, nil)
	s.cache.On("SaveToken", mock.Anything).Run(func(a mock.Arguments){
		t, ok := a.Get(0).(*oidc.Token)
		s.Require().True(ok)

		s.Assert().Equal(token.AccessToken, t.AccessToken)
		s.Assert().Empty(t.IDToken)
		s.Assert().Equal(token.RefreshToken, t.RefreshToken)
	}).Return(nil)

	resetDone := true
	clear := s.oidcSource.clearIDToken(func(){resetDone = true})
	s.Require().NoError(clear())

	s.cache.AssertExpectations(s.T())
}<|endoftext|>


<commit_before>package zap

import (
	"crypto/tls"
	"encoding/json"
	"log"
	"net"
	"net/http"
	"time"

	"github.com/moomerman/zap/cert"
	"github.com/puma/puma-dev/dev/launch"
	"github.com/unrolled/render"
	"golang.org/x/net/http2"
)

// Server holds the state for the HTTP and HTTPS servers
type Server struct {
	http  *http.Server
	https *http.Server
}

// NewServer starts the HTTP and HTTPS proxy servers
func NewServer() *Server {
	httpsMux := http.NewServeMux()
	// TODO: don't handle these requests unless localhost request (eg. not via ngrok)
	httpsMux.HandleFunc("/zap/api/log", logAPIHandler())
	httpsMux.HandleFunc("/zap/api/state", stateAPIHandler())
	httpsMux.HandleFunc("/zap/api/apps", appsAPIHandler())
	httpsMux.HandleFunc("/zap/log", logHandler())
	httpsMux.HandleFunc("/zap/restart", restartHandler())
	httpsMux.HandleFunc("/zap", statusHandler())
	httpsMux.HandleFunc("/", appHandler())

	httpMux := http.NewServeMux()
	httpMux.HandleFunc("/", appHandler())

	http := startHTTP(httpMux)
	https := startHTTPS(httpsMux)

	return &Server{
		http:  http,
		https: https,
	}
}

var renderer = render.New(render.Options{
	Layout:     "layout",
	Asset:      Asset,
	AssetNames: AssetNames,
	Extensions: []string{".html"},
})

// ServeTLS starts the HTTPS server
func (s *Server) ServeTLS(bind string) {
	if bind == "SocketTLS" {
		listeners, err := launch.SocketListeners(bind)
		if err != nil {
			log.Fatal("unable to get launchd socket listener", err)
		}

		s.https.Serve(tls.NewListener(listeners[0], s.https.TLSConfig))
	} else {
		listener, err := tls.Listen("tcp", bind, s.https.TLSConfig)
		if err != nil {
			log.Fatal("unable to create tls listener", err)
		}
		s.https.Serve(listener)
	}
}

// Serve starts the HTTP server
func (s *Server) Serve(bind string) {
	if bind == "Socket" {
		listeners, err := launch.SocketListeners(bind)
		if err != nil {
			log.Fatal("unable to get launchd socket listener", err)
		}
		s.http.Serve(listeners[0])
	} else {
		listener, err := net.Listen("tcp", bind)
		if err != nil {
			log.Fatal("unable to create listener", err)
		}
		s.http.Serve(listener)
	}
}

func startHTTPS(handler http.Handler) *http.Server {
	cache, err := cert.NewCache()
	if err != nil {
		log.Fatal("unable to create new cert cache", err)
	}

	tlsConfig := &tls.Config{
		GetCertificate: cache.GetCertificate,
	}

	server := &http.Server{
		Handler:   handler,
		TLSConfig: tlsConfig,
	}
	http2.ConfigureServer(server, nil)

	return server
}

func startHTTP(handler http.Handler) *http.Server {
	return &http.Server{
		Handler: handler,
	}
}

func appHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app, err := findAppForHost(r.Host)
		if err != nil {
			renderer.HTML(w, http.StatusBadGateway, "502", "App Not Found")
			return
		}

		switch app.Status() {
		case "running":
			app.ServeHTTP(w, r)
		default:
			renderer.HTML(w, http.StatusAccepted, "app", app)
		}
	}
}

func statusHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app, err := findAppForHost(r.Host)
		if err != nil {
			renderer.HTML(w, http.StatusBadGateway, "502", "App Not Found")
			return
		}

		renderer.HTML(w, http.StatusOK, "app", app)
	}
}

func logHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app, err := findAppForHost(r.Host)
		if err != nil {
			renderer.HTML(w, http.StatusBadGateway, "502", "App Not Found")
			return
		}

		renderer.HTML(w, http.StatusOK, "log", app)
	}
}

func restartHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {

		app, err := findAppForHost(r.Host)
		if err != nil {
			renderer.HTML(w, http.StatusBadGateway, "502", "App Not Found")
			return
		}

		if err := app.Restart(); err != nil {
			log.Println("[app]", app.Config.Host, "internal server error", err)
			http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
		}

		http.Redirect(w, r, "/zap", http.StatusTemporaryRedirect)
	}
}

func logAPIHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app, err := findAppForHost(r.Host)
		if err != nil {
			http.Error(w, "502 App Not Found", http.StatusBadGateway)
			return
		}

		app.WriteLog(w)
	}
}

func stateAPIHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app, err := findAppForHost(r.Host)
		if err != nil {
			renderer.HTML(w, http.StatusBadGateway, "502", "App Not Found")
			return
		}

		content, err := json.MarshalIndent(map[string]interface{}{
			"app":    app,
			"uptime": time.Since(app.Started).String(),
			"status": app.Status(),
		}, "", "  ")

		w.Header().Set("Content-Type", "application/json")
		w.Write(content)
	}
}

func appsAPIHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {

		content, err := json.MarshalIndent(map[string]interface{}{
			"apps": apps,
		}, "", "  ")
		if err != nil {
			log.Println("[app] internal server error", err)
			http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.Write(content)
	}
}
<commit_msg>Dry up app loading logic in handlers<commit_after>package zap

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"log"
	"net"
	"net/http"
	"time"

	"github.com/moomerman/zap/cert"
	"github.com/puma/puma-dev/dev/launch"
	"github.com/unrolled/render"
	"golang.org/x/net/http2"
)

// Server holds the state for the HTTP and HTTPS servers
type Server struct {
	http  *http.Server
	https *http.Server
}

// NewServer starts the HTTP and HTTPS proxy servers
func NewServer() *Server {
	httpsMux := http.NewServeMux()
	// TODO: don't handle these requests unless localhost request (eg. not via ngrok)
	httpsMux.HandleFunc("/zap/api/log", findAppHandler(logAPIHandler()))
	httpsMux.HandleFunc("/zap/api/state", findAppHandler(stateAPIHandler()))
	httpsMux.HandleFunc("/zap/api/apps", appsAPIHandler())
	httpsMux.HandleFunc("/zap/log", findAppHandler(logHandler()))
	httpsMux.HandleFunc("/zap/restart", findAppHandler(restartHandler()))
	httpsMux.HandleFunc("/zap", findAppHandler(statusHandler()))
	httpsMux.HandleFunc("/", findAppHandler(appHandler()))

	httpMux := http.NewServeMux()
	httpMux.HandleFunc("/", findAppHandler(appHandler()))

	http := startHTTP(httpMux)
	https := startHTTPS(httpsMux)

	return &Server{
		http:  http,
		https: https,
	}
}

var renderer = render.New(render.Options{
	Layout:     "layout",
	Asset:      Asset,
	AssetNames: AssetNames,
	Extensions: []string{".html"},
})

// ServeTLS starts the HTTPS server
func (s *Server) ServeTLS(bind string) {
	if bind == "SocketTLS" {
		listeners, err := launch.SocketListeners(bind)
		if err != nil {
			log.Fatal("unable to get launchd socket listener", err)
		}

		s.https.Serve(tls.NewListener(listeners[0], s.https.TLSConfig))
	} else {
		listener, err := tls.Listen("tcp", bind, s.https.TLSConfig)
		if err != nil {
			log.Fatal("unable to create tls listener", err)
		}
		s.https.Serve(listener)
	}
}

// Serve starts the HTTP server
func (s *Server) Serve(bind string) {
	if bind == "Socket" {
		listeners, err := launch.SocketListeners(bind)
		if err != nil {
			log.Fatal("unable to get launchd socket listener", err)
		}
		s.http.Serve(listeners[0])
	} else {
		listener, err := net.Listen("tcp", bind)
		if err != nil {
			log.Fatal("unable to create listener", err)
		}
		s.http.Serve(listener)
	}
}

func startHTTPS(handler http.Handler) *http.Server {
	cache, err := cert.NewCache()
	if err != nil {
		log.Fatal("unable to create new cert cache", err)
	}

	tlsConfig := &tls.Config{
		GetCertificate: cache.GetCertificate,
	}

	server := &http.Server{
		Handler:   handler,
		TLSConfig: tlsConfig,
	}
	http2.ConfigureServer(server, nil)

	return server
}

func startHTTP(handler http.Handler) *http.Server {
	return &http.Server{
		Handler: handler,
	}
}

type contextKey string

var appKey contextKey = "app"

func findAppHandler(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		app, err := findAppForHost(r.Host)
		if err != nil {
			renderer.HTML(w, http.StatusBadGateway, "502", "App Not Found")
			return
		}
		ctx := context.WithValue(r.Context(), appKey, app)
		next.ServeHTTP(w, r.WithContext(ctx))
	}
}

func appHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app := r.Context().Value(appKey).(*app)

		switch app.Status() {
		case "running":
			app.ServeHTTP(w, r)
		default:
			renderer.HTML(w, http.StatusAccepted, "app", app)
		}
	}
}

func statusHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app := r.Context().Value(appKey).(*app)
		renderer.HTML(w, http.StatusOK, "app", app)
	}
}

func logHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app := r.Context().Value(appKey).(*app)
		renderer.HTML(w, http.StatusOK, "log", app)
	}
}

func restartHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app := r.Context().Value(appKey).(*app)

		if err := app.Restart(); err != nil {
			log.Println("[app]", app.Config.Host, "internal server error", err)
			http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
		}

		http.Redirect(w, r, "/zap", http.StatusTemporaryRedirect)
	}
}

func logAPIHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app := r.Context().Value(appKey).(*app)

		app.WriteLog(w)
	}
}

func stateAPIHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		app := r.Context().Value(appKey).(*app)

		content, _ := json.MarshalIndent(map[string]interface{}{
			"app":    app,
			"uptime": time.Since(app.Started).String(),
			"status": app.Status(),
		}, "", "  ")

		w.Header().Set("Content-Type", "application/json")
		w.Write(content)
	}
}

func appsAPIHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {

		content, err := json.MarshalIndent(map[string]interface{}{
			"apps": apps,
		}, "", "  ")
		if err != nil {
			log.Println("[app] internal server error", err)
			http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.Write(content)
	}
}
<|endoftext|>


<commit_before>package repository

import (
	"github.com/lib/pq"
	"github.com/syou6162/go-active-learning/lib/model"
)

func (r *repository) UpdateRelatedExamples(related model.RelatedExamples) error {
	if _, err := r.db.Exec(`DELETE FROM related_example WHERE example_id = $1;`, related.ExampleId); err != nil {
		return err
	}
	if _, err := r.db.Exec(`INSERT INTO related_example (example_id, related_example_id) VALUES ($1, unnest(cast($2 AS INT[])));`, related.ExampleId, pq.Array(related.RelatedExampleIds)); err != nil {
		return err
	}
	return nil
}

func (r *repository) FindRelatedExamples(e *model.Example) (*model.RelatedExamples, error) {
	related := &model.RelatedExamples{ExampleId: e.Id}
	items := make([]int, 0)
	query := `SELECT example_id FROM related_example WHERE example_id = $1;`
	err := r.db.Select(&items, query, e.Id)
	if err != nil {
		return nil, err
	}
	related.RelatedExampleIds = items
	return related, nil
}
<commit_msg>選択するカラムを間違えていた<commit_after>package repository

import (
	"github.com/lib/pq"
	"github.com/syou6162/go-active-learning/lib/model"
)

func (r *repository) UpdateRelatedExamples(related model.RelatedExamples) error {
	if _, err := r.db.Exec(`DELETE FROM related_example WHERE example_id = $1;`, related.ExampleId); err != nil {
		return err
	}
	if _, err := r.db.Exec(`INSERT INTO related_example (example_id, related_example_id) VALUES ($1, unnest(cast($2 AS INT[])));`, related.ExampleId, pq.Array(related.RelatedExampleIds)); err != nil {
		return err
	}
	return nil
}

func (r *repository) FindRelatedExamples(e *model.Example) (*model.RelatedExamples, error) {
	related := &model.RelatedExamples{ExampleId: e.Id}
	items := make([]int, 0)
	query := `SELECT related_example_id FROM related_example WHERE example_id = $1;`
	err := r.db.Select(&items, query, e.Id)
	if err != nil {
		return nil, err
	}
	related.RelatedExampleIds = items
	return related, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build go1.3

package lxdclient

import (
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/juju/errors"
	"github.com/lxc/lxd/shared"
)

const (
	tempPrefix = "juju-lxd-"

	pemBlockTypeCert = "CERTIFICATE"
	pemBlockTypeKey  = "RSA PRIVATE KEY"
)

// Cert holds the information for a single certificate a client
// may use to connect to a remote server.
type Cert struct {
	// Name is the name that LXD will use for the cert.
	Name string

	// CertPEM is the PEM-encoded x.509 cert.
	CertPEM []byte

	// KeyPEM is the PEM-encoded x.509 private key.
	KeyPEM []byte
}

// NewCert creates a new Cert for the given cert and key.
func NewCert(certPEM, keyPEM []byte) *Cert {
	return &Cert{
		CertPEM: certPEM,
		KeyPEM:  keyPEM,
	}
}

// Validate ensures that the cert is valid.
func (cert Cert) Validate() error {
	if len(cert.CertPEM) == 0 {
		return errors.NotValidf("missing cert PEM")
	}
	if len(cert.KeyPEM) == 0 {
		return errors.NotValidf("missing key PEM")
	}

	// TODO(ericsnow) Ensure cert and key are valid?

	return nil
}

// WriteCertPEM writes the cert's x.509 PEM data to the given writer.
func (cert Cert) WriteCertPEM(out io.Writer) error {
	if _, err := out.Write(cert.CertPEM); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// WriteKeytPEM writes the key's x.509 PEM data to the given writer.
func (cert Cert) WriteKeyPEM(out io.Writer) error {
	if _, err := out.Write(cert.KeyPEM); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// Fingerprint returns the cert's LXD fingerprint.
func (cert Cert) Fingerprint() (string, error) {
	// See: https://github.com/lxc/lxd/blob/master/lxd/certificates.go
	x509Cert, err := cert.X509()
	if err != nil {
		return "", errors.Trace(err)
	}
	data := sha256.Sum256(x509Cert.Raw)
	return fmt.Sprintf("%x", data), nil
}

// X509 returns the x.509 certificate.
func (cert Cert) X509() (*x509.Certificate, error) {
	block, _ := pem.Decode(cert.CertPEM)
	if block == nil {
		return nil, errors.Errorf("invalid cert PEM (%d bytes)", len(cert.CertPEM))
	}

	x509Cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, errors.Trace(err)
	}

	return x509Cert, nil
}

func genCertAndKey() ([]byte, []byte, error) {
	// See GenCert() in:
	//  https://github.com/lxc/lxd/blob/master/shared/cert.go
	// TODO(ericsnow) Split up GenCert so it is more re-usable.
	tempdir, err := ioutil.TempDir("", tempPrefix)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	defer os.RemoveAll(tempdir)
	certFile := filepath.Join(tempdir, configCertFile)
	keyFile := filepath.Join(tempdir, configKeyFile)
	if err := shared.GenCert(certFile, keyFile); err != nil {
		return nil, nil, errors.Trace(err)
	}

	certPEM, err := ioutil.ReadFile(certFile)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	keyPEM, err := ioutil.ReadFile(keyFile)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	return certPEM, keyPEM, nil
}
<commit_msg>Add Cert.SetDefaults().<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build go1.3

package lxdclient

import (
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/juju/errors"
	"github.com/lxc/lxd/shared"
)

const (
	tempPrefix = "juju-lxd-"

	pemBlockTypeCert = "CERTIFICATE"
	pemBlockTypeKey  = "RSA PRIVATE KEY"

	// TODO(ericsnow) Is this the right default?
	certDefaultName = configCertFile
)

// Cert holds the information for a single certificate a client
// may use to connect to a remote server.
type Cert struct {
	// Name is the name that LXD will use for the cert.
	Name string

	// CertPEM is the PEM-encoded x.509 cert.
	CertPEM []byte

	// KeyPEM is the PEM-encoded x.509 private key.
	KeyPEM []byte
}

// NewCert creates a new Cert for the given cert and key.
func NewCert(certPEM, keyPEM []byte) *Cert {
	return &Cert{
		CertPEM: certPEM,
		KeyPEM:  keyPEM,
	}
}

// SetDefaults updates a copy of the remote with default values
// where needed.
func (cert Cert) SetDefaults() (*Cert, error) {
	if cert.Name == "" {
		cert.Name = certDefaultName
	}

	// TODO(ericsnow) populate cert/key (use genCertAndKey)?

	return &cert, nil
}

// Validate ensures that the cert is valid.
func (cert Cert) Validate() error {
	if len(cert.CertPEM) == 0 {
		return errors.NotValidf("missing cert PEM")
	}
	if len(cert.KeyPEM) == 0 {
		return errors.NotValidf("missing key PEM")
	}

	// TODO(ericsnow) Ensure cert and key are valid?

	return nil
}

// WriteCertPEM writes the cert's x.509 PEM data to the given writer.
func (cert Cert) WriteCertPEM(out io.Writer) error {
	if _, err := out.Write(cert.CertPEM); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// WriteKeytPEM writes the key's x.509 PEM data to the given writer.
func (cert Cert) WriteKeyPEM(out io.Writer) error {
	if _, err := out.Write(cert.KeyPEM); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// Fingerprint returns the cert's LXD fingerprint.
func (cert Cert) Fingerprint() (string, error) {
	// See: https://github.com/lxc/lxd/blob/master/lxd/certificates.go
	x509Cert, err := cert.X509()
	if err != nil {
		return "", errors.Trace(err)
	}
	data := sha256.Sum256(x509Cert.Raw)
	return fmt.Sprintf("%x", data), nil
}

// X509 returns the x.509 certificate.
func (cert Cert) X509() (*x509.Certificate, error) {
	block, _ := pem.Decode(cert.CertPEM)
	if block == nil {
		return nil, errors.Errorf("invalid cert PEM (%d bytes)", len(cert.CertPEM))
	}

	x509Cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, errors.Trace(err)
	}

	return x509Cert, nil
}

func genCertAndKey() ([]byte, []byte, error) {
	// See GenCert() in:
	//  https://github.com/lxc/lxd/blob/master/shared/cert.go
	// TODO(ericsnow) Split up GenCert so it is more re-usable.
	tempdir, err := ioutil.TempDir("", tempPrefix)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	defer os.RemoveAll(tempdir)
	certFile := filepath.Join(tempdir, configCertFile)
	keyFile := filepath.Join(tempdir, configKeyFile)
	if err := shared.GenCert(certFile, keyFile); err != nil {
		return nil, nil, errors.Trace(err)
	}

	certPEM, err := ioutil.ReadFile(certFile)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	keyPEM, err := ioutil.ReadFile(keyFile)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	return certPEM, keyPEM, nil
}
<|endoftext|>


<commit_before>package ams

import (
	"net/http"
	"testing"

	"context"
	"net/http/httptest"
	"reflect"
	_ "reflect"
)

func TestClient_GetMediaProcessors(t *testing.T) {
	expected := []MediaProcessor{
		{
			ID:          "sample-id1",
			Name:        "Sample Media Processor",
			Description: "Sample Media Processor, for test",
			SKU:         "Stock Keeping Unit?",
			Vendor:      "Sample Vendor",
			Version:     "1.0.0",
		},
		{
			ID:          "sample-id2",
			Name:        "Standard Media Processor",
			Description: "Standard Media Processor, for test",
			SKU:         "Stock Keeping Unit??",
			Vendor:      "Sample Inc.",
			Version:     "1.1.0",
		},
		{
			ID:          "sample-id3",
			Name:        "Simple Media Processor",
			Description: "Simple Media Processor, for test",
			SKU:         "Stock Keeping Unit???",
			Vendor:      "Simple Vendor",
			Version:     "2.0.0",
		},
	}
	m := http.NewServeMux()
	m.HandleFunc("/MediaProcessors",
		testJSONHandler(t, http.MethodGet, false, http.StatusOK, testWrapValue(expected)),
	)

	s := httptest.NewServer(m)
	defer s.Close()

	client := testClient(t, s.URL)

	actual, err := client.GetMediaProcessors(context.TODO())
	if err != nil {
		t.Fatal(err)
	}

	if !reflect.DeepEqual(actual, expected) {
		t.Errorf("unexpected media processors. expected: %#v, actual: %#v", expected, actual)
	}
}
<commit_msg>chore: remove ignore import<commit_after>package ams

import (
	"net/http"
	"testing"

	"context"
	"net/http/httptest"
	"reflect"
)

func TestClient_GetMediaProcessors(t *testing.T) {
	expected := []MediaProcessor{
		{
			ID:          "sample-id1",
			Name:        "Sample Media Processor",
			Description: "Sample Media Processor, for test",
			SKU:         "Stock Keeping Unit?",
			Vendor:      "Sample Vendor",
			Version:     "1.0.0",
		},
		{
			ID:          "sample-id2",
			Name:        "Standard Media Processor",
			Description: "Standard Media Processor, for test",
			SKU:         "Stock Keeping Unit??",
			Vendor:      "Sample Inc.",
			Version:     "1.1.0",
		},
		{
			ID:          "sample-id3",
			Name:        "Simple Media Processor",
			Description: "Simple Media Processor, for test",
			SKU:         "Stock Keeping Unit???",
			Vendor:      "Simple Vendor",
			Version:     "2.0.0",
		},
	}
	m := http.NewServeMux()
	m.HandleFunc("/MediaProcessors",
		testJSONHandler(t, http.MethodGet, false, http.StatusOK, testWrapValue(expected)),
	)

	s := httptest.NewServer(m)
	defer s.Close()

	client := testClient(t, s.URL)

	actual, err := client.GetMediaProcessors(context.TODO())
	if err != nil {
		t.Fatal(err)
	}

	if !reflect.DeepEqual(actual, expected) {
		t.Errorf("unexpected media processors. expected: %#v, actual: %#v", expected, actual)
	}
}
<|endoftext|>


<commit_before>// instances provides for running a bunch of skiaserve instances.
package instances

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"regexp"
	"sort"
	"sync"
	"time"

	"github.com/google/uuid"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	// MAX_INSTANCES is the max number of concurrent skiaserve instances we
	// support in the hosted environment.
	MAX_INSTANCES = 200

	// MIN_POOL_SIZE is the number of available spots in the pool that should be
	// maintained at all times. Instances will be culled by oldest lastUsed time
	// until this size is reached.
	MIN_POOL_SIZE = 10

	// START_PORT Is the beginning of the range of ports the skiaserve instances
	// will communicate on.
	START_PORT = 30000

	// START_WAIT_NUM poll the newly started skiaserve this many times before giving up.
	START_WAIT_NUM = 50

	// START_WAIT_PERIOD poll the newly started skiaserve this often.
	START_WAIT_PERIOD = 100 * time.Millisecond

	// EXIT_WAIT_PERIOD is the time to wait for the instance to exit.
	EXIT_WAIT_PERIOD = 2 * time.Second

	// SKIASERVE is the full path to the skiaserve executable.
	SKIASERVE = "/usr/local/bin/skiaserve"
)

var (
	// instancePrefixRe is used to strip out the instance uuid.
	instancePrefixRe = regexp.MustCompile("^/([A-Fa-f0-9-]+)(/.*)")

	runningInstances = metrics2.GetCounter("running_instances", nil)
)

// NewInstanceID creates a new id for an instance of skiaserve.
func NewInstanceID() string {
	return uuid.New().String()
}

// instance represents a single skiaserve instance, which may or may not
// be running. It is used in Instances.
type instance struct {
	// proxy is the proxy connection to talk to the running skiaserve.
	proxy *httputil.ReverseProxy

	// port is the port that skiaserve is listening on.
	port int

	display int

	// uuid is the login id of the uuid this skiaserve is running for.
	uuid string // "" means this isn't running.

	// lastUsed is the time the skiaserve instance last processed a request.
	lastUsed time.Time

	// started is the time that the skiaserve instance was started. Will be used
	// later when we give hosted users the ability to see if their skiaserve is
	// out of date.
	started time.Time

	process exec.Process
}

// Start a single instance of skiaserve running at the given port.
//
func (c *instance) Start(uuid string) (<-chan error, error) {
	runCmd := &exec.Command{
		Name:      "xvfb-run",
		Args:      []string{"--server-args", "-screen 0 1280x1024x24", "--server-num", fmt.Sprintf("%d", c.port), SKIASERVE, "--port", fmt.Sprintf("%d", c.port), "--source", "https://debugger-assets.skia.org", "--hosted"},
		Env:       []string{fmt.Sprintf("DISPLAY=:%d", c.display)},
		LogStderr: true,
		LogStdout: true,
	}
	process, exitChan, err := exec.RunIndefinitely(runCmd)
	if err == nil {
		c.process = process
		c.uuid = uuid
	}
	return exitChan, err
}

// Stop the instance from running.
func (c *instance) Stop() {
	if err := c.process.Kill(); err != nil {
		sklog.Errorf("Error trying to kill instance: %s", err)
	}
}

// Sort so oldest with a non-empty uuid are first.
type instanceSlice []*instance

func (p instanceSlice) Len() int { return len(p) }
func (p instanceSlice) Less(i, j int) bool {
	// if p[i].uuid == "" XOR p[j].uuid == ""
	if (p[i].uuid == "") != (p[j].uuid == "") {
		return p[i].uuid > p[j].uuid
	} else {
		return p[i].lastUsed.Before(p[j].lastUsed)
	}
}
func (p instanceSlice) Swap(i, j int) { p[i], p[j] = p[j], p[i] }

// Instances is used to control a number of skiaserve instances all running
// at the same time.
//
// Instances implements http.handler, which reverse proxies incoming requests
// to the right backend.
//
// TODO(jcgregorio) Need to add memory limits to instance.
type Instances struct {
	// pool is the list of potential running skiaserve instances. We only start
	// them on demand.
	pool []*instance

	// instances is a map from uuid to a running skiaserve.
	instances map[string]*instance

	// mutex protects access to pool and instances.
	mutex sync.Mutex
}

func New() *Instances {
	s := &Instances{
		pool:      []*instance{},
		instances: map[string]*instance{},
	}
	for i := 0; i < MAX_INSTANCES; i++ {
		port := START_PORT + i
		proxyurl := fmt.Sprintf("http://localhost:%d", port)
		u, err := url.Parse(proxyurl)
		if err != nil {
			sklog.Errorf("failed to parse url %q: %s", proxyurl, err)
		}
		c := &instance{
			port:    port,
			display: i + 100,
			proxy:   httputil.NewSingleHostReverseProxy(u),
		}
		s.pool = append(s.pool, c)
	}

	// Start a culling process that runs every minute that keeps
	// some open spots in the pool.
	go func() {
		for range time.Tick(time.Minute) {
			s.cull()
		}
	}()

	return s
}

// cull terminates the oldest running instances to make room for more.
func (s *Instances) cull() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	sort.Sort(instanceSlice(s.pool))
	cull := s.pool[len(s.pool)-MIN_POOL_SIZE:]
	for _, co := range cull {
		if co.uuid != "" {
			co.Stop()
		}
	}
}

// getInstanceFromPool finds the first open instance in the pool.
func (s *Instances) getInstanceFromPool(uuid string) *instance {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	// Find first open instance in the pool.
	var co *instance = nil
	for _, c := range s.pool {
		if c.uuid == "" {
			c.uuid = uuid
			co = c
			break
		}
	}
	if co != nil {
		s.instances[uuid] = co
	}
	return co
}

// startInstance starts skiaserve running in a instance for the given uuid.
//
// It waits until skiaserve responds to an HTTP request before returning.
//
// The actual instance for the uuid is determined by looking at the prefix of
// the URL.Path, i.e. /2..F/foo will be directed to instance 2..F for the given uuid
// and skiaserve will be sent the URL.Path "/foo", i.e. with the instance
// number prefix stripped.
func (s *Instances) startInstance(uuid string) error {
	co := s.getInstanceFromPool(uuid)
	if co == nil {
		return fmt.Errorf("Could not start an instance.")
	}
	runningInstances := metrics2.GetCounter("running_instances", nil)
	runningInstances.Inc(1)
	co.started = time.Now()
	exitChan, err := co.Start(uuid)
	if err != nil {
		return fmt.Errorf("Failed to start instance at port %d: %s", co.port, err)
	}

	go func() {
		<-exitChan
		// Remove the entry for this instance now that it has exited.
		s.mutex.Lock()
		defer s.mutex.Unlock()
		delete(s.instances, uuid)
		runningInstances.Dec(1)
		co.uuid = ""
	}()

	// Poll the port until we get a response.
	url := fmt.Sprintf("http://localhost:%d", co.port)
	var resp *http.Response
	client := httputils.NewTimeoutClient()
	for i := 0; i < START_WAIT_NUM; i++ {
		resp, err = client.Get(url)
		if resp != nil && resp.Body != nil {
			util.Close(resp.Body)
		}
		if err == nil {
			break
		}
		time.Sleep(START_WAIT_PERIOD)
	}
	if err != nil {
		return fmt.Errorf("Started instance but skiaserve never responded: %s", err)
	}

	return nil
}

// getInstance returns the instance for the given uuid, or nil if there isn't
// one for that uuid.
func (s *Instances) getInstance(uuid string) *instance {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.instances[uuid]
}

// setLastUsed set the lastUsed timestamp for an instance.
func (s *Instances) setLastUsed(uuid string) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.instances[uuid].lastUsed = time.Now()
}

// ServeHTTP implements the http.Handler interface by proxying the requests to
// the correct instance based on the uuid.
//
// All requests are routed to the instance, with the exception of
// /instanceStatus and /instanceNew which will be handled by 'co' itself.
func (s *Instances) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/instanceNew" && r.Method == "POST" {
		// TODO(jcgregorio) Add gorilla.csrf protection.
		http.Redirect(w, r, fmt.Sprintf("/%s/", NewInstanceID()), 303)
		return
	}
	instanceID := ""
	// Strip off the uuid from the URL.
	parts := instancePrefixRe.FindStringSubmatch(r.URL.Path)
	sklog.Infof("Parts: %v", parts)
	if len(parts) == 3 {
		instanceID = parts[1]
		r.URL.Path = parts[2]
	} else {
		httputils.ReportError(w, r, fmt.Errorf("Invalid URL %q", r.URL.Path), "Not a valid URL.")
		return
	}

	co := s.getInstance(instanceID)
	if co == nil {
		// If no instance then start one up.
		if err := s.startInstance(instanceID); err != nil {
			httputils.ReportError(w, r, err, "Failed to start new instance.")
			return
		}
		co = s.getInstance(instanceID)
		if co == nil {
			httputils.ReportError(w, r, fmt.Errorf("Failed to start instance %q", instanceID), "Started instance, but then couldn't find it.")
			return
		}
	}

	if r.URL.Path == "/instanceStatus" {
		if r.Method == "GET" {
			// A GET to /instanceStatus will return the instance info, i.e. how long it's been running.
			enc := json.NewEncoder(w)
			if err := enc.Encode(
				struct {
					Started int64 `json:"started"`
				}{
					Started: co.started.Unix(),
				},
			); err != nil {
				httputils.ReportError(w, r, err, "Failed to serialize response.")
			}
			return
		} else if r.Method == "POST" {
			co.Stop()
			time.Sleep(EXIT_WAIT_PERIOD)
			s.mutex.Lock()
			defer s.mutex.Unlock()
			// Remove the entry for this instance now that it has exited.
			delete(s.instances, instanceID)
			http.Redirect(w, r, fmt.Sprintf("/%s/", instanceID), 303)
		}
		return
	}

	// Proxy.
	sklog.Infof("Proxying request: %s %s", r.URL, instanceID)
	// If the request is a POST and we are at a non-zero instanceNum then pass in
	// a recording response.  If the response is a 303 then we return a 303 with
	// the correct location URL, otherwise we return the response verbatim.
	if r.Method == "POST" {
		// TODO(jcgregorio) If this is an uploaded SKP then store it in GCS
		// to make restarts less painful.
		rw := httptest.NewRecorder()
		co.proxy.ServeHTTP(rw, r)
		if rw.Code == 303 {
			http.Redirect(w, r, fmt.Sprintf("/%s/", instanceID), 303)
		} else {
			for k, values := range rw.HeaderMap {
				for _, v := range values {
					w.Header().Set(k, v)
				}
			}
			if _, err := w.Write(rw.Body.Bytes()); err != nil {
				sklog.Errorf("Failed proxying a recorded response: %s", err)
			}
		}
	} else {
		co.proxy.ServeHTTP(w, r)
	}

	// Update lastUsed.
	s.setLastUsed(instanceID)
}

// StopAll stops all running instances.
func (s *Instances) StopAll() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	for _, co := range s.instances {
		sklog.Infof("Stopping instance for uuid %q on port %d", co.uuid, co.port)
		co.Stop()
	}
}

type InstanceInfo struct {
	ID     string        `json:"id"`
	UUID   string        `json:"uuid"`
	Uptime time.Duration `json:"uptime"`
	Port   int           `json:"port"`
}

type InstanceInfoSlice []*InstanceInfo

func (p InstanceInfoSlice) Len() int           { return len(p) }
func (p InstanceInfoSlice) Less(i, j int) bool { return p[i].ID < p[j].ID }
func (p InstanceInfoSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

func (s *Instances) DescribeAll() []*InstanceInfo {
	info := []*InstanceInfo{}
	s.mutex.Lock()
	defer s.mutex.Unlock()
	for id, co := range s.instances {
		info = append(info, &InstanceInfo{
			ID:     id,
			UUID:   co.uuid,
			Uptime: time.Now().Sub(co.started),
			Port:   co.port,
		})
	}
	sort.Sort(InstanceInfoSlice(info))
	return info
}
<commit_msg>debugger - Stop spamming the error logs.<commit_after>// instances provides for running a bunch of skiaserve instances.
package instances

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"regexp"
	"sort"
	"sync"
	"time"

	"github.com/google/uuid"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	// MAX_INSTANCES is the max number of concurrent skiaserve instances we
	// support in the hosted environment.
	MAX_INSTANCES = 200

	// MIN_POOL_SIZE is the number of available spots in the pool that should be
	// maintained at all times. Instances will be culled by oldest lastUsed time
	// until this size is reached.
	MIN_POOL_SIZE = 10

	// START_PORT Is the beginning of the range of ports the skiaserve instances
	// will communicate on.
	START_PORT = 30000

	// START_WAIT_NUM poll the newly started skiaserve this many times before giving up.
	START_WAIT_NUM = 50

	// START_WAIT_PERIOD poll the newly started skiaserve this often.
	START_WAIT_PERIOD = 100 * time.Millisecond

	// EXIT_WAIT_PERIOD is the time to wait for the instance to exit.
	EXIT_WAIT_PERIOD = 2 * time.Second

	// SKIASERVE is the full path to the skiaserve executable.
	SKIASERVE = "/usr/local/bin/skiaserve"
)

var (
	// instancePrefixRe is used to strip out the instance uuid.
	instancePrefixRe = regexp.MustCompile("^/([A-Fa-f0-9-]+)(/.*)")

	runningInstances = metrics2.GetCounter("running_instances", nil)
)

// NewInstanceID creates a new id for an instance of skiaserve.
func NewInstanceID() string {
	return uuid.New().String()
}

// instance represents a single skiaserve instance, which may or may not
// be running. It is used in Instances.
type instance struct {
	// proxy is the proxy connection to talk to the running skiaserve.
	proxy *httputil.ReverseProxy

	// port is the port that skiaserve is listening on.
	port int

	display int

	// uuid is the login id of the uuid this skiaserve is running for.
	uuid string // "" means this isn't running.

	// lastUsed is the time the skiaserve instance last processed a request.
	lastUsed time.Time

	// started is the time that the skiaserve instance was started. Will be used
	// later when we give hosted users the ability to see if their skiaserve is
	// out of date.
	started time.Time

	process exec.Process
}

// Start a single instance of skiaserve running at the given port.
//
func (c *instance) Start(uuid string) (<-chan error, error) {
	runCmd := &exec.Command{
		Name:      "xvfb-run",
		Args:      []string{"--server-args", "-screen 0 1280x1024x24", "--server-num", fmt.Sprintf("%d", c.port), SKIASERVE, "--port", fmt.Sprintf("%d", c.port), "--source", "https://debugger-assets.skia.org", "--hosted"},
		Env:       []string{fmt.Sprintf("DISPLAY=:%d", c.display)},
		LogStdout: true,
	}
	process, exitChan, err := exec.RunIndefinitely(runCmd)
	if err == nil {
		c.process = process
		c.uuid = uuid
	}
	return exitChan, err
}

// Stop the instance from running.
func (c *instance) Stop() {
	if err := c.process.Kill(); err != nil {
		sklog.Errorf("Error trying to kill instance: %s", err)
	}
}

// Sort so oldest with a non-empty uuid are first.
type instanceSlice []*instance

func (p instanceSlice) Len() int { return len(p) }
func (p instanceSlice) Less(i, j int) bool {
	// if p[i].uuid == "" XOR p[j].uuid == ""
	if (p[i].uuid == "") != (p[j].uuid == "") {
		return p[i].uuid > p[j].uuid
	} else {
		return p[i].lastUsed.Before(p[j].lastUsed)
	}
}
func (p instanceSlice) Swap(i, j int) { p[i], p[j] = p[j], p[i] }

// Instances is used to control a number of skiaserve instances all running
// at the same time.
//
// Instances implements http.handler, which reverse proxies incoming requests
// to the right backend.
//
// TODO(jcgregorio) Need to add memory limits to instance.
type Instances struct {
	// pool is the list of potential running skiaserve instances. We only start
	// them on demand.
	pool []*instance

	// instances is a map from uuid to a running skiaserve.
	instances map[string]*instance

	// mutex protects access to pool and instances.
	mutex sync.Mutex
}

func New() *Instances {
	s := &Instances{
		pool:      []*instance{},
		instances: map[string]*instance{},
	}
	for i := 0; i < MAX_INSTANCES; i++ {
		port := START_PORT + i
		proxyurl := fmt.Sprintf("http://localhost:%d", port)
		u, err := url.Parse(proxyurl)
		if err != nil {
			sklog.Errorf("failed to parse url %q: %s", proxyurl, err)
		}
		c := &instance{
			port:    port,
			display: i + 100,
			proxy:   httputil.NewSingleHostReverseProxy(u),
		}
		s.pool = append(s.pool, c)
	}

	// Start a culling process that runs every minute that keeps
	// some open spots in the pool.
	go func() {
		for range time.Tick(time.Minute) {
			s.cull()
		}
	}()

	return s
}

// cull terminates the oldest running instances to make room for more.
func (s *Instances) cull() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	sort.Sort(instanceSlice(s.pool))
	cull := s.pool[len(s.pool)-MIN_POOL_SIZE:]
	for _, co := range cull {
		if co.uuid != "" {
			co.Stop()
		}
	}
}

// getInstanceFromPool finds the first open instance in the pool.
func (s *Instances) getInstanceFromPool(uuid string) *instance {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	// Find first open instance in the pool.
	var co *instance = nil
	for _, c := range s.pool {
		if c.uuid == "" {
			c.uuid = uuid
			co = c
			break
		}
	}
	if co != nil {
		s.instances[uuid] = co
	}
	return co
}

// startInstance starts skiaserve running in a instance for the given uuid.
//
// It waits until skiaserve responds to an HTTP request before returning.
//
// The actual instance for the uuid is determined by looking at the prefix of
// the URL.Path, i.e. /2..F/foo will be directed to instance 2..F for the given uuid
// and skiaserve will be sent the URL.Path "/foo", i.e. with the instance
// number prefix stripped.
func (s *Instances) startInstance(uuid string) error {
	co := s.getInstanceFromPool(uuid)
	if co == nil {
		return fmt.Errorf("Could not start an instance.")
	}
	runningInstances := metrics2.GetCounter("running_instances", nil)
	runningInstances.Inc(1)
	co.started = time.Now()
	exitChan, err := co.Start(uuid)
	if err != nil {
		return fmt.Errorf("Failed to start instance at port %d: %s", co.port, err)
	}

	go func() {
		<-exitChan
		// Remove the entry for this instance now that it has exited.
		s.mutex.Lock()
		defer s.mutex.Unlock()
		delete(s.instances, uuid)
		runningInstances.Dec(1)
		co.uuid = ""
	}()

	// Poll the port until we get a response.
	url := fmt.Sprintf("http://localhost:%d", co.port)
	var resp *http.Response
	client := httputils.NewTimeoutClient()
	for i := 0; i < START_WAIT_NUM; i++ {
		resp, err = client.Get(url)
		if resp != nil && resp.Body != nil {
			util.Close(resp.Body)
		}
		if err == nil {
			break
		}
		time.Sleep(START_WAIT_PERIOD)
	}
	if err != nil {
		return fmt.Errorf("Started instance but skiaserve never responded: %s", err)
	}

	return nil
}

// getInstance returns the instance for the given uuid, or nil if there isn't
// one for that uuid.
func (s *Instances) getInstance(uuid string) *instance {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	return s.instances[uuid]
}

// setLastUsed set the lastUsed timestamp for an instance.
func (s *Instances) setLastUsed(uuid string) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.instances[uuid].lastUsed = time.Now()
}

// ServeHTTP implements the http.Handler interface by proxying the requests to
// the correct instance based on the uuid.
//
// All requests are routed to the instance, with the exception of
// /instanceStatus and /instanceNew which will be handled by 'co' itself.
func (s *Instances) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/instanceNew" && r.Method == "POST" {
		// TODO(jcgregorio) Add gorilla.csrf protection.
		http.Redirect(w, r, fmt.Sprintf("/%s/", NewInstanceID()), 303)
		return
	}
	instanceID := ""
	// Strip off the uuid from the URL.
	parts := instancePrefixRe.FindStringSubmatch(r.URL.Path)
	sklog.Infof("Parts: %v", parts)
	if len(parts) == 3 {
		instanceID = parts[1]
		r.URL.Path = parts[2]
	} else {
		httputils.ReportError(w, r, fmt.Errorf("Invalid URL %q", r.URL.Path), "Not a valid URL.")
		return
	}

	co := s.getInstance(instanceID)
	if co == nil {
		// If no instance then start one up.
		if err := s.startInstance(instanceID); err != nil {
			httputils.ReportError(w, r, err, "Failed to start new instance.")
			return
		}
		co = s.getInstance(instanceID)
		if co == nil {
			httputils.ReportError(w, r, fmt.Errorf("Failed to start instance %q", instanceID), "Started instance, but then couldn't find it.")
			return
		}
	}

	if r.URL.Path == "/instanceStatus" {
		if r.Method == "GET" {
			// A GET to /instanceStatus will return the instance info, i.e. how long it's been running.
			enc := json.NewEncoder(w)
			if err := enc.Encode(
				struct {
					Started int64 `json:"started"`
				}{
					Started: co.started.Unix(),
				},
			); err != nil {
				httputils.ReportError(w, r, err, "Failed to serialize response.")
			}
			return
		} else if r.Method == "POST" {
			co.Stop()
			time.Sleep(EXIT_WAIT_PERIOD)
			s.mutex.Lock()
			defer s.mutex.Unlock()
			// Remove the entry for this instance now that it has exited.
			delete(s.instances, instanceID)
			http.Redirect(w, r, fmt.Sprintf("/%s/", instanceID), 303)
		}
		return
	}

	// Proxy.
	sklog.Infof("Proxying request: %s %s", r.URL, instanceID)
	// If the request is a POST and we are at a non-zero instanceNum then pass in
	// a recording response.  If the response is a 303 then we return a 303 with
	// the correct location URL, otherwise we return the response verbatim.
	if r.Method == "POST" {
		// TODO(jcgregorio) If this is an uploaded SKP then store it in GCS
		// to make restarts less painful.
		rw := httptest.NewRecorder()
		co.proxy.ServeHTTP(rw, r)
		if rw.Code == 303 {
			http.Redirect(w, r, fmt.Sprintf("/%s/", instanceID), 303)
		} else {
			for k, values := range rw.HeaderMap {
				for _, v := range values {
					w.Header().Set(k, v)
				}
			}
			if _, err := w.Write(rw.Body.Bytes()); err != nil {
				sklog.Errorf("Failed proxying a recorded response: %s", err)
			}
		}
	} else {
		co.proxy.ServeHTTP(w, r)
	}

	// Update lastUsed.
	s.setLastUsed(instanceID)
}

// StopAll stops all running instances.
func (s *Instances) StopAll() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	for _, co := range s.instances {
		sklog.Infof("Stopping instance for uuid %q on port %d", co.uuid, co.port)
		co.Stop()
	}
}

type InstanceInfo struct {
	ID     string        `json:"id"`
	UUID   string        `json:"uuid"`
	Uptime time.Duration `json:"uptime"`
	Port   int           `json:"port"`
}

type InstanceInfoSlice []*InstanceInfo

func (p InstanceInfoSlice) Len() int           { return len(p) }
func (p InstanceInfoSlice) Less(i, j int) bool { return p[i].ID < p[j].ID }
func (p InstanceInfoSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

func (s *Instances) DescribeAll() []*InstanceInfo {
	info := []*InstanceInfo{}
	s.mutex.Lock()
	defer s.mutex.Unlock()
	for id, co := range s.instances {
		info = append(info, &InstanceInfo{
			ID:     id,
			UUID:   co.uuid,
			Uptime: time.Now().Sub(co.started),
			Port:   co.port,
		})
	}
	sort.Sort(InstanceInfoSlice(info))
	return info
}
<|endoftext|>


<commit_before>package email

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"log"
	"mime/multipart"
	"net"
	"net/mail"
	"net/smtp"
	"net/textproto"
	"strings"
	"time"

	"github.com/dockercn/docker-bucket/models"
)

const (
	maxLineLength = 76
)

//默认5分钟发一次
func MailService() {
	go func() {
		for {
			//加载模板列表去到prefix的集合，再去prefix对应的message
			tmpl := new(models.TemplateHtml)
			tmpls := tmpl.Query()
			for i, _ := range tmpls {
				msg := new(models.Message)
				msgs := msg.Query(tmpls[i].Prefix)
				for j, _ := range msgs {
					fmt.Printf("%#v\n", msgs[j])
					mailServer := new(models.MailServer)
					server := mailServer.Query(msgs[j].Host)
					isSend, _ := Send(server[0], msgs[j])
					if isSend {
						msgs[j].Update()
					}
				}
			}
			time.Sleep(5 * time.Minute)
		}
	}()
}

//调用发送邮件服务 返回值bool表示是否发送  error表示如果发送是成功还是失败
func Send(mailServer *models.MailServer, msg *models.Message) (bool, error) {
	//判断邮件状态  I待发送 A 已经发送 X 发送失败    发送失败超过10次，停止发送
	if msg.Status == "A" {
		return false, nil
	} else if msg.Status == "X" && msg.Count > 10 {
		return false, nil
	}
	auth := smtp.PlainAuth("", mailServer.User, mailServer.Password, mailServer.Host)
	to := make([]string, 0, len(msg.Cc)+len(msg.Bcc)+1)
	to = append(append(append(to, msg.To), msg.Cc...), msg.Bcc...)
	// Check to make sure there is at least one recipient and one "From" address
	if msg.From == "" || len(to) == 0 {
		return false, errors.New("Must specify at least one From address and one To address")
	}
	from, err := mail.ParseAddress(msg.From)
	if err != nil {
		return false, err
	}
	raw, err := msg2bytes(msg)
	if err != nil {
		return false, err
	}
	//err = SendMail(mailServer.Host+":"+strconv.Itoa(mailServer.Port), auth, from.Address, to, raw)
	address := fmt.Sprintf("%s:%d", mailServer.Host, 465)
	err = sendMailUsingTLS(address, auth, from.Address, to, raw)
	//判断发送是否成功
	msg.Count = msg.Count + 1
	if err != nil {
		msg.Status = "X"
		return true, err
	}
	msg.Status = "A"
	return true, nil
}

func msg2bytes(e *models.Message) ([]byte, error) {
	e.Headers = textproto.MIMEHeader{}
	buff := &bytes.Buffer{}
	w := multipart.NewWriter(buff)
	// Set the appropriate headers (overwriting any conflicts)
	// Leave out Bcc (only included in envelope headers)
	e.Headers.Set("To", e.To)
	if e.Cc != nil {
		e.Headers.Set("Cc", strings.Join(e.Cc, ","))
	}
	e.Headers.Set("From", e.From)
	e.Headers.Set("Subject", e.Subject)
	e.Headers.Set("MIME-Version", "1.0")
	e.Headers.Set("Content-Type", fmt.Sprintf("multipart/mixed;\r\n boundary=%s\r\n", w.Boundary()))
	// Write the envelope headers (including any custom headers)
	if err := headerToBytes(buff, e.Headers); err != nil {
		return nil, fmt.Errorf("Failed to render message headers: %s", err)
	}
	// Start the multipart/mixed part
	fmt.Fprintf(buff, "--%s\r\n", w.Boundary())
	header := textproto.MIMEHeader{}
	// Check to see if there is a Text or HTML field
	if e.Type != "" {
		subWriter := multipart.NewWriter(buff)
		// Create the multipart alternative part
		header.Set("Content-Type", fmt.Sprintf("multipart/alternative;\r\n boundary=%s\r\n", subWriter.Boundary()))
		// Write the header
		if err := headerToBytes(buff, header); err != nil {
			return nil, fmt.Errorf("Failed to render multipart message headers: %s", err)
		}
		// Create the body sections
		if e.Type == "text" {
			header.Set("Content-Type", fmt.Sprintf("text/plain; charset=UTF-8"))
			header.Set("Content-Transfer-Encoding", "quoted-printable")
			if _, err := subWriter.CreatePart(header); err != nil {
				return nil, err
			}
			// Write the text
			if err := quotePrintEncode(buff, e.Body); err != nil {
				return nil, err
			}
		}
		if e.Type == "html" {
			header.Set("Content-Type", fmt.Sprintf("text/html; charset=UTF-8"))
			header.Set("Content-Transfer-Encoding", "quoted-printable")
			if _, err := subWriter.CreatePart(header); err != nil {
				return nil, err
			}
			// Write the text
			if err := quotePrintEncode(buff, e.Body); err != nil {
				return nil, err
			}
		}
		if err := subWriter.Close(); err != nil {
			return nil, err
		}
	}
	// Create attachment part, if necessary
	for _, a := range e.Attachments {
		ap, err := w.CreatePart(a.Header)
		if err != nil {
			return nil, err
		}
		// Write the base64Wrapped content to the part
		base64Wrap(ap, a.Content)
	}
	if err := w.Close(); err != nil {
		return nil, err
	}
	return buff.Bytes(), nil
}

// headerToBytes enumerates the key and values in the header, and writes the results to the IO Writer
func headerToBytes(w io.Writer, t textproto.MIMEHeader) error {
	for k, v := range t {
		// Write the header key
		_, err := fmt.Fprintf(w, "%s:", k)
		if err != nil {
			return err
		}
		// Write each value in the header
		for _, c := range v {
			_, err := fmt.Fprintf(w, " %s\r\n", c)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func quotePrintEncode(w io.Writer, s string) error {
	var buf [3]byte
	mc := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		// We're assuming Unix style text formats as input (LF line break), and
		// quoted-printble uses CRLF line breaks. (Literal CRs will become
		// "=0D", but probably shouldn't be there to begin with!)
		if c == '\n' {
			io.WriteString(w, "\r\n")
			mc = 0
			continue
		}
		var nextOut []byte
		if isPrintable(c) {
			nextOut = append(buf[:0], c)
		} else {
			nextOut = buf[:]
			qpEscape(nextOut, c)
		}
		// Add a soft line break if the next (encoded) byte would push this line
		// to or past the limit.
		if mc+len(nextOut) >= maxLineLength {
			if _, err := io.WriteString(w, "=\r\n"); err != nil {
				return err
			}
			mc = 0
		}
		if _, err := w.Write(nextOut); err != nil {
			return err
		}
		mc += len(nextOut)
	}
	// No trailing end-of-line?? Soft line break, then. TODO: is this sane?
	if mc > 0 {
		io.WriteString(w, "=\r\n")
	}
	return nil
}

// base64Wrap encodes the attachment content, and wraps it according to RFC 2045 standards (every 76 chars)
// The output is then written to the specified io.Writer
func base64Wrap(w io.Writer, b []byte) {
	// 57 raw bytes per 76-byte base64 line.
	const maxRaw = 57
	// Buffer for each line, including trailing CRLF.
	var buffer [maxLineLength + len("\r\n")]byte
	copy(buffer[maxLineLength:], "\r\n")
	// Process raw chunks until there's no longer enough to fill a line.
	for len(b) >= maxRaw {
		base64.StdEncoding.Encode(buffer[:], b[:maxRaw])
		w.Write(buffer[:])
		b = b[maxRaw:]
	}
	// Handle the last chunk of bytes.
	if len(b) > 0 {
		out := buffer[:base64.StdEncoding.EncodedLen(len(b))]
		base64.StdEncoding.Encode(out, b)
		out = append(out, "\r\n"...)
		w.Write(out)
	}
}

// qpEscape is a helper function for quotePrintEncode which escapes a
// non-printable byte. Expects len(dest) == 3.
func qpEscape(dest []byte, c byte) {
	const nums = "0123456789ABCDEF"
	dest[0] = '='
	dest[1] = nums[(c&0xf0)>>4]
	dest[2] = nums[(c & 0xf)]
}

// isPrintable returns true if the rune given is "printable" according to RFC 2045, false otherwise
func isPrintable(c byte) bool {
	return (c >= '!' && c <= '<') || (c >= '>' && c <= '~') || (c == ' ' || c == '\n' || c == '\t')
}

func sendMailUsingTLS(addr string, auth smtp.Auth, from string, to []string, msg []byte) (err error) {
	//create smtp client
	c, err := dial(addr)
	if err != nil {
		log.Println("Create smpt client error:", err)
		return err
	}
	defer c.Close()
	if auth != nil {
		if ok, _ := c.Extension("AUTH"); ok {
			if err = c.Auth(auth); err != nil {
				log.Println("Error during AUTH", err)
				return err
			}
		}
	}
	if err = c.Mail(from); err != nil {
		return err
	}
	for _, addr := range to {
		if err = c.Rcpt(addr); err != nil {
			return err
		}
	}
	w, err := c.Data()
	if err != nil {
		return err
	}
	_, err = w.Write(msg)
	if err != nil {
		return err
	}
	err = w.Close()
	if err != nil {
		return err
	}
	return c.Quit()
}

func dial(addr string) (*smtp.Client, error) {
	conn, err := tls.Dial("tcp", addr, nil)
	if err != nil {
		log.Println("Dialing Error:", err)
		return nil, err
	}
	//分解主机端口字符串
	host, _, _ := net.SplitHostPort(addr)
	return smtp.NewClient(conn, host)
}
<commit_msg>add smtp server check<commit_after>package email

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"log"
	"mime/multipart"
	"net"
	"net/mail"
	"net/smtp"
	"net/textproto"
	"strings"
	"time"

	"github.com/dockercn/docker-bucket/models"
)

const (
	maxLineLength = 76
)

//默认5分钟发一次
func MailService() {
	checkService()
	go func() {
		for {
			//加载模板列表去到prefix的集合，再去prefix对应的message
			tmpl := new(models.TemplateHtml)
			tmpls := tmpl.Query()
			for i, _ := range tmpls {
				msg := new(models.Message)
				msgs := msg.Query(tmpls[i].Prefix)
				for j, _ := range msgs {
					fmt.Printf("%#v\n", msgs[j])
					mailServer := new(models.MailServer)
					server := mailServer.Query(msgs[j].Host)
					isSend, _ := Send(server[0], msgs[j])
					if isSend {
						msgs[j].Update()
					}
				}
			}
			time.Sleep(5 * time.Minute)
		}
	}()
}

func checkService() {
	mailServer := new(models.MailServer)
	servers := mailServer.Query()
	if len(servers) == 0 {
		log.Println("数据库中未设置任何smtp服务器,无法发送任何邮件")
		return
	}
	log.Println("邮件服务器设置检查完成，邮件服务正常启动")
}

//调用发送邮件服务 返回值bool表示是否发送  error表示如果发送是成功还是失败
func Send(mailServer *models.MailServer, msg *models.Message) (bool, error) {
	//判断邮件状态  I待发送 A 已经发送 X 发送失败    发送失败超过10次，停止发送
	if msg.Status == "A" {
		return false, nil
	} else if msg.Status == "X" && msg.Count > 10 {
		return false, nil
	}
	auth := smtp.PlainAuth("", mailServer.User, mailServer.Password, mailServer.Host)
	to := make([]string, 0, len(msg.Cc)+len(msg.Bcc)+1)
	to = append(append(append(to, msg.To), msg.Cc...), msg.Bcc...)
	// Check to make sure there is at least one recipient and one "From" address
	if msg.From == "" || len(to) == 0 {
		return false, errors.New("Must specify at least one From address and one To address")
	}
	from, err := mail.ParseAddress(msg.From)
	if err != nil {
		return false, err
	}
	raw, err := msg2bytes(msg)
	if err != nil {
		return false, err
	}
	//err = SendMail(mailServer.Host+":"+strconv.Itoa(mailServer.Port), auth, from.Address, to, raw)
	address := fmt.Sprintf("%s:%d", mailServer.Host, 465)
	err = sendMailUsingTLS(address, auth, from.Address, to, raw)
	//判断发送是否成功
	msg.Count = msg.Count + 1
	if err != nil {
		msg.Status = "X"
		return true, err
	}
	msg.Status = "A"
	return true, nil
}

func msg2bytes(e *models.Message) ([]byte, error) {
	e.Headers = textproto.MIMEHeader{}
	buff := &bytes.Buffer{}
	w := multipart.NewWriter(buff)
	// Set the appropriate headers (overwriting any conflicts)
	// Leave out Bcc (only included in envelope headers)
	e.Headers.Set("To", e.To)
	if e.Cc != nil {
		e.Headers.Set("Cc", strings.Join(e.Cc, ","))
	}
	e.Headers.Set("From", e.From)
	e.Headers.Set("Subject", e.Subject)
	e.Headers.Set("MIME-Version", "1.0")
	e.Headers.Set("Content-Type", fmt.Sprintf("multipart/mixed;\r\n boundary=%s\r\n", w.Boundary()))
	// Write the envelope headers (including any custom headers)
	if err := headerToBytes(buff, e.Headers); err != nil {
		return nil, fmt.Errorf("Failed to render message headers: %s", err)
	}
	// Start the multipart/mixed part
	fmt.Fprintf(buff, "--%s\r\n", w.Boundary())
	header := textproto.MIMEHeader{}
	// Check to see if there is a Text or HTML field
	if e.Type != "" {
		subWriter := multipart.NewWriter(buff)
		// Create the multipart alternative part
		header.Set("Content-Type", fmt.Sprintf("multipart/alternative;\r\n boundary=%s\r\n", subWriter.Boundary()))
		// Write the header
		if err := headerToBytes(buff, header); err != nil {
			return nil, fmt.Errorf("Failed to render multipart message headers: %s", err)
		}
		// Create the body sections
		if e.Type == "text" {
			header.Set("Content-Type", fmt.Sprintf("text/plain; charset=UTF-8"))
			header.Set("Content-Transfer-Encoding", "quoted-printable")
			if _, err := subWriter.CreatePart(header); err != nil {
				return nil, err
			}
			// Write the text
			if err := quotePrintEncode(buff, e.Body); err != nil {
				return nil, err
			}
		}
		if e.Type == "html" {
			header.Set("Content-Type", fmt.Sprintf("text/html; charset=UTF-8"))
			header.Set("Content-Transfer-Encoding", "quoted-printable")
			if _, err := subWriter.CreatePart(header); err != nil {
				return nil, err
			}
			// Write the text
			if err := quotePrintEncode(buff, e.Body); err != nil {
				return nil, err
			}
		}
		if err := subWriter.Close(); err != nil {
			return nil, err
		}
	}
	// Create attachment part, if necessary
	for _, a := range e.Attachments {
		ap, err := w.CreatePart(a.Header)
		if err != nil {
			return nil, err
		}
		// Write the base64Wrapped content to the part
		base64Wrap(ap, a.Content)
	}
	if err := w.Close(); err != nil {
		return nil, err
	}
	return buff.Bytes(), nil
}

// headerToBytes enumerates the key and values in the header, and writes the results to the IO Writer
func headerToBytes(w io.Writer, t textproto.MIMEHeader) error {
	for k, v := range t {
		// Write the header key
		_, err := fmt.Fprintf(w, "%s:", k)
		if err != nil {
			return err
		}
		// Write each value in the header
		for _, c := range v {
			_, err := fmt.Fprintf(w, " %s\r\n", c)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func quotePrintEncode(w io.Writer, s string) error {
	var buf [3]byte
	mc := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		// We're assuming Unix style text formats as input (LF line break), and
		// quoted-printble uses CRLF line breaks. (Literal CRs will become
		// "=0D", but probably shouldn't be there to begin with!)
		if c == '\n' {
			io.WriteString(w, "\r\n")
			mc = 0
			continue
		}
		var nextOut []byte
		if isPrintable(c) {
			nextOut = append(buf[:0], c)
		} else {
			nextOut = buf[:]
			qpEscape(nextOut, c)
		}
		// Add a soft line break if the next (encoded) byte would push this line
		// to or past the limit.
		if mc+len(nextOut) >= maxLineLength {
			if _, err := io.WriteString(w, "=\r\n"); err != nil {
				return err
			}
			mc = 0
		}
		if _, err := w.Write(nextOut); err != nil {
			return err
		}
		mc += len(nextOut)
	}
	// No trailing end-of-line?? Soft line break, then. TODO: is this sane?
	if mc > 0 {
		io.WriteString(w, "=\r\n")
	}
	return nil
}

// base64Wrap encodes the attachment content, and wraps it according to RFC 2045 standards (every 76 chars)
// The output is then written to the specified io.Writer
func base64Wrap(w io.Writer, b []byte) {
	// 57 raw bytes per 76-byte base64 line.
	const maxRaw = 57
	// Buffer for each line, including trailing CRLF.
	var buffer [maxLineLength + len("\r\n")]byte
	copy(buffer[maxLineLength:], "\r\n")
	// Process raw chunks until there's no longer enough to fill a line.
	for len(b) >= maxRaw {
		base64.StdEncoding.Encode(buffer[:], b[:maxRaw])
		w.Write(buffer[:])
		b = b[maxRaw:]
	}
	// Handle the last chunk of bytes.
	if len(b) > 0 {
		out := buffer[:base64.StdEncoding.EncodedLen(len(b))]
		base64.StdEncoding.Encode(out, b)
		out = append(out, "\r\n"...)
		w.Write(out)
	}
}

// qpEscape is a helper function for quotePrintEncode which escapes a
// non-printable byte. Expects len(dest) == 3.
func qpEscape(dest []byte, c byte) {
	const nums = "0123456789ABCDEF"
	dest[0] = '='
	dest[1] = nums[(c&0xf0)>>4]
	dest[2] = nums[(c & 0xf)]
}

// isPrintable returns true if the rune given is "printable" according to RFC 2045, false otherwise
func isPrintable(c byte) bool {
	return (c >= '!' && c <= '<') || (c >= '>' && c <= '~') || (c == ' ' || c == '\n' || c == '\t')
}

func sendMailUsingTLS(addr string, auth smtp.Auth, from string, to []string, msg []byte) (err error) {
	//create smtp client
	c, err := dial(addr)
	if err != nil {
		log.Println("Create smpt client error:", err)
		return err
	}
	defer c.Close()
	if auth != nil {
		if ok, _ := c.Extension("AUTH"); ok {
			if err = c.Auth(auth); err != nil {
				log.Println("Error during AUTH", err)
				return err
			}
		}
	}
	if err = c.Mail(from); err != nil {
		return err
	}
	for _, addr := range to {
		if err = c.Rcpt(addr); err != nil {
			return err
		}
	}
	w, err := c.Data()
	if err != nil {
		return err
	}
	_, err = w.Write(msg)
	if err != nil {
		return err
	}
	err = w.Close()
	if err != nil {
		return err
	}
	return c.Quit()
}

func dial(addr string) (*smtp.Client, error) {
	conn, err := tls.Dial("tcp", addr, nil)
	if err != nil {
		log.Println("Dialing Error:", err)
		return nil, err
	}
	//分解主机端口字符串
	host, _, _ := net.SplitHostPort(addr)
	return smtp.NewClient(conn, host)
}
<|endoftext|>