<commit_before>package agent

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"

	"github.com/docker/docker/pkg/ioutils"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/ugorji/go/codec"
)

var (
	allocIDNotPresentErr  = fmt.Errorf("must provide a valid alloc id")
	fileNameNotPresentErr = fmt.Errorf("must provide a file name")
	taskNotPresentErr     = fmt.Errorf("must provide task name")
	logTypeNotPresentErr  = fmt.Errorf("must provide log type (stdout/stderr)")
	clientNotRunning      = fmt.Errorf("node is not running a Nomad Client")
	invalidOrigin         = fmt.Errorf("origin must be start or end")
)

func (s *HTTPServer) FsRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	path := strings.TrimPrefix(req.URL.Path, "/v1/client/fs/")
	switch {
	case strings.HasPrefix(path, "ls/"):
		return s.DirectoryListRequest(resp, req)
	case strings.HasPrefix(path, "stat/"):
		return s.FileStatRequest(resp, req)
	case strings.HasPrefix(path, "readat/"):
		return s.FileReadAtRequest(resp, req)
	case strings.HasPrefix(path, "cat/"):
		return s.FileCatRequest(resp, req)
	case strings.HasPrefix(path, "stream/"):
		return s.Stream(resp, req)
	case strings.HasPrefix(path, "logs/"):
		return s.Logs(resp, req)
	default:
		return nil, CodedError(404, ErrInvalidMethod)
	}
}

func (s *HTTPServer) DirectoryListRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string

	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/ls/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}
	if path = req.URL.Query().Get("path"); path == "" {
		path = "/"
	}

	// Create the request
	args := &cstructs.FsListRequest{
		AllocID: allocID,
		Path:    path,
	}
	s.parse(resp, req, &args.QueryOptions.Region, &args.QueryOptions)

	// Make the RPC
	localClient, remoteClient, localServer := s.rpcHandlerForAlloc(allocID)

	var reply cstructs.FsListResponse
	var rpcErr error
	if localClient {
		rpcErr = s.agent.Client().ClientRPC("FileSystem.List", &args, &reply)
	} else if remoteClient {
		rpcErr = s.agent.Client().RPC("FileSystem.List", &args, &reply)
	} else if localServer {
		rpcErr = s.agent.Server().RPC("FileSystem.List", &args, &reply)
	}

	if rpcErr != nil {
		if structs.IsErrNoNodeConn(rpcErr) || structs.IsErrUnknownAllocation(rpcErr) {
			rpcErr = CodedError(404, rpcErr.Error())
		}

		return nil, rpcErr
	}

	return reply.Files, nil
}

func (s *HTTPServer) FileStatRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string
	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/stat/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}
	if path = req.URL.Query().Get("path"); path == "" {
		return nil, fileNameNotPresentErr
	}

	// Create the request
	args := &cstructs.FsStatRequest{
		AllocID: allocID,
		Path:    path,
	}
	s.parse(resp, req, &args.QueryOptions.Region, &args.QueryOptions)

	// Make the RPC
	localClient, remoteClient, localServer := s.rpcHandlerForAlloc(allocID)

	var reply cstructs.FsStatResponse
	var rpcErr error
	if localClient {
		rpcErr = s.agent.Client().ClientRPC("FileSystem.Stat", &args, &reply)
	} else if remoteClient {
		rpcErr = s.agent.Client().RPC("FileSystem.Stat", &args, &reply)
	} else if localServer {
		rpcErr = s.agent.Server().RPC("FileSystem.Stat", &args, &reply)
	}

	if rpcErr != nil {
		if structs.IsErrNoNodeConn(rpcErr) || structs.IsErrUnknownAllocation(rpcErr) {
			rpcErr = CodedError(404, rpcErr.Error())
		}

		return nil, rpcErr
	}

	return reply.Info, nil
}

func (s *HTTPServer) FileReadAtRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string
	var offset, limit int64
	var err error

	q := req.URL.Query()

	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/readat/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}
	if path = q.Get("path"); path == "" {
		return nil, fileNameNotPresentErr
	}

	if offset, err = strconv.ParseInt(q.Get("offset"), 10, 64); err != nil {
		return nil, fmt.Errorf("error parsing offset: %v", err)
	}

	// Parse the limit
	if limitStr := q.Get("limit"); limitStr != "" {
		if limit, err = strconv.ParseInt(limitStr, 10, 64); err != nil {
			return nil, fmt.Errorf("error parsing limit: %v", err)
		}
	}

	// Create the request arguments
	fsReq := &cstructs.FsStreamRequest{
		AllocID:   allocID,
		Path:      path,
		Offset:    offset,
		Origin:    "start",
		Limit:     limit,
		PlainText: true,
	}
	s.parse(resp, req, &fsReq.QueryOptions.Region, &fsReq.QueryOptions)

	// Make the request
	return s.fsStreamImpl(resp, req, "FileSystem.Stream", fsReq, fsReq.AllocID)
}

func (s *HTTPServer) FileCatRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string

	q := req.URL.Query()

	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/cat/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}
	if path = q.Get("path"); path == "" {
		return nil, fileNameNotPresentErr
	}

	// Create the request arguments
	fsReq := &cstructs.FsStreamRequest{
		AllocID:   allocID,
		Path:      path,
		Origin:    "start",
		PlainText: true,
	}
	s.parse(resp, req, &fsReq.QueryOptions.Region, &fsReq.QueryOptions)

	// Make the request
	return s.fsStreamImpl(resp, req, "FileSystem.Stream", fsReq, fsReq.AllocID)
}

// Stream streams the content of a file blocking on EOF.
// The parameters are:
// * path: path to file to stream.
// * offset: The offset to start streaming data at, defaults to zero.
// * origin: Either "start" or "end" and defines from where the offset is
//           applied. Defaults to "start".
func (s *HTTPServer) Stream(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string

	q := req.URL.Query()

	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/stream/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}

	if path = q.Get("path"); path == "" {
		return nil, fileNameNotPresentErr
	}

	var offset int64
	offsetString := q.Get("offset")
	if offsetString != "" {
		var err error
		if offset, err = strconv.ParseInt(offsetString, 10, 64); err != nil {
			return nil, fmt.Errorf("error parsing offset: %v", err)
		}
	}

	origin := q.Get("origin")
	switch origin {
	case "start", "end":
	case "":
		origin = "start"
	default:
		return nil, invalidOrigin
	}

	// Create the request arguments
	fsReq := &cstructs.FsStreamRequest{
		AllocID: allocID,
		Path:    path,
		Origin:  origin,
		Offset:  offset,
		Follow:  true,
	}
	s.parse(resp, req, &fsReq.QueryOptions.Region, &fsReq.QueryOptions)

	// Make the request
	return s.fsStreamImpl(resp, req, "FileSystem.Stream", fsReq, fsReq.AllocID)
}

// Logs streams the content of a log blocking on EOF. The parameters are:
// * task: task name to stream logs for.
// * type: stdout/stderr to stream.
// * follow: A boolean of whether to follow the logs.
// * offset: The offset to start streaming data at, defaults to zero.
// * origin: Either "start" or "end" and defines from where the offset is
//           applied. Defaults to "start".
func (s *HTTPServer) Logs(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, task, logType string
	var plain, follow bool
	var err error

	q := req.URL.Query()
	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/logs/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}

	if task = q.Get("task"); task == "" {
		return nil, taskNotPresentErr
	}

	if followStr := q.Get("follow"); followStr != "" {
		if follow, err = strconv.ParseBool(followStr); err != nil {
			return nil, fmt.Errorf("Failed to parse follow field to boolean: %v", err)
		}
	}

	if plainStr := q.Get("plain"); plainStr != "" {
		if plain, err = strconv.ParseBool(plainStr); err != nil {
			return nil, fmt.Errorf("Failed to parse plain field to boolean: %v", err)
		}
	}

	logType = q.Get("type")
	switch logType {
	case "stdout", "stderr":
	default:
		return nil, logTypeNotPresentErr
	}

	var offset int64
	offsetString := q.Get("offset")
	if offsetString != "" {
		var err error
		if offset, err = strconv.ParseInt(offsetString, 10, 64); err != nil {
			return nil, fmt.Errorf("error parsing offset: %v", err)
		}
	}

	origin := q.Get("origin")
	switch origin {
	case "start", "end":
	case "":
		origin = "start"
	default:
		return nil, invalidOrigin
	}

	// Create the request arguments
	fsReq := &cstructs.FsLogsRequest{
		AllocID:   allocID,
		Task:      task,
		LogType:   logType,
		Offset:    offset,
		Origin:    origin,
		PlainText: plain,
		Follow:    follow,
	}
	s.parse(resp, req, &fsReq.QueryOptions.Region, &fsReq.QueryOptions)

	// Make the request
	return s.fsStreamImpl(resp, req, "FileSystem.Logs", fsReq, fsReq.AllocID)
}

// fsStreamImpl is used to make a streaming filesystem call that serializes the
// args and then expects a stream of StreamErrWrapper results where the payload
// is copied to the response body.
func (s *HTTPServer) fsStreamImpl(resp http.ResponseWriter,
	req *http.Request, method string, args interface{}, allocID string) (interface{}, error) {

	// Get the correct handler
	localClient, remoteClient, localServer := s.rpcHandlerForAlloc(allocID)
	var handler structs.StreamingRpcHandler
	var handlerErr error
	if localClient {
		handler, handlerErr = s.agent.Client().StreamingRpcHandler(method)
	} else if remoteClient {
		handler, handlerErr = s.agent.Client().RemoteStreamingRpcHandler(method)
	} else if localServer {
		handler, handlerErr = s.agent.Server().StreamingRpcHandler(method)
	}

	if handlerErr != nil {
		return nil, CodedError(500, handlerErr.Error())
	}

	// Create a pipe connecting the (possibly remote) handler to the http response
	httpPipe, handlerPipe := net.Pipe()
	decoder := codec.NewDecoder(httpPipe, structs.MsgpackHandle)
	encoder := codec.NewEncoder(httpPipe, structs.MsgpackHandle)

	// Create a goroutine that closes the pipe if the connection closes.
	ctx, cancel := context.WithCancel(req.Context())
	go func() {
		<-ctx.Done()
		httpPipe.Close()
	}()

	// Create an output that gets flushed on every write
	output := ioutils.NewWriteFlusher(resp)

	// Create a channel that decodes the results
	errCh := make(chan HTTPCodedError)
	go func() {
		defer cancel()

		// Send the request
		if err := encoder.Encode(args); err != nil {
			errCh <- CodedError(500, err.Error())
			return
		}
		encoder.Reset(httpPipe)

		for {
			select {
			case <-ctx.Done():
				errCh <- nil
				return
			default:
			}

			var res cstructs.StreamErrWrapper
			if err := decoder.Decode(&res); err != nil {
				errCh <- CodedError(500, err.Error())
				return
			}
			decoder.Reset(httpPipe)

			if err := res.Error; err != nil {
				if err.Code != nil {
					errCh <- CodedError(int(*err.Code), err.Error())
					return
				}
			}

			if _, err := io.Copy(output, bytes.NewBuffer(res.Payload)); err != nil {
				errCh <- CodedError(500, err.Error())
				return
			}
		}
	}()

	handler(handlerPipe)
	cancel()
	codedErr := <-errCh
	if codedErr != nil &&
		(codedErr == io.EOF ||
			strings.Contains(codedErr.Error(), "closed") ||
			strings.Contains(codedErr.Error(), "EOF")) {
		codedErr = nil
	}
	return nil, codedErr
}
<commit_msg>client: use a bytes.Reader for reading a []byte<commit_after>package agent

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"

	"github.com/docker/docker/pkg/ioutils"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/ugorji/go/codec"
)

var (
	allocIDNotPresentErr  = fmt.Errorf("must provide a valid alloc id")
	fileNameNotPresentErr = fmt.Errorf("must provide a file name")
	taskNotPresentErr     = fmt.Errorf("must provide task name")
	logTypeNotPresentErr  = fmt.Errorf("must provide log type (stdout/stderr)")
	clientNotRunning      = fmt.Errorf("node is not running a Nomad Client")
	invalidOrigin         = fmt.Errorf("origin must be start or end")
)

func (s *HTTPServer) FsRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	path := strings.TrimPrefix(req.URL.Path, "/v1/client/fs/")
	switch {
	case strings.HasPrefix(path, "ls/"):
		return s.DirectoryListRequest(resp, req)
	case strings.HasPrefix(path, "stat/"):
		return s.FileStatRequest(resp, req)
	case strings.HasPrefix(path, "readat/"):
		return s.FileReadAtRequest(resp, req)
	case strings.HasPrefix(path, "cat/"):
		return s.FileCatRequest(resp, req)
	case strings.HasPrefix(path, "stream/"):
		return s.Stream(resp, req)
	case strings.HasPrefix(path, "logs/"):
		return s.Logs(resp, req)
	default:
		return nil, CodedError(404, ErrInvalidMethod)
	}
}

func (s *HTTPServer) DirectoryListRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string

	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/ls/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}
	if path = req.URL.Query().Get("path"); path == "" {
		path = "/"
	}

	// Create the request
	args := &cstructs.FsListRequest{
		AllocID: allocID,
		Path:    path,
	}
	s.parse(resp, req, &args.QueryOptions.Region, &args.QueryOptions)

	// Make the RPC
	localClient, remoteClient, localServer := s.rpcHandlerForAlloc(allocID)

	var reply cstructs.FsListResponse
	var rpcErr error
	if localClient {
		rpcErr = s.agent.Client().ClientRPC("FileSystem.List", &args, &reply)
	} else if remoteClient {
		rpcErr = s.agent.Client().RPC("FileSystem.List", &args, &reply)
	} else if localServer {
		rpcErr = s.agent.Server().RPC("FileSystem.List", &args, &reply)
	}

	if rpcErr != nil {
		if structs.IsErrNoNodeConn(rpcErr) || structs.IsErrUnknownAllocation(rpcErr) {
			rpcErr = CodedError(404, rpcErr.Error())
		}

		return nil, rpcErr
	}

	return reply.Files, nil
}

func (s *HTTPServer) FileStatRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string
	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/stat/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}
	if path = req.URL.Query().Get("path"); path == "" {
		return nil, fileNameNotPresentErr
	}

	// Create the request
	args := &cstructs.FsStatRequest{
		AllocID: allocID,
		Path:    path,
	}
	s.parse(resp, req, &args.QueryOptions.Region, &args.QueryOptions)

	// Make the RPC
	localClient, remoteClient, localServer := s.rpcHandlerForAlloc(allocID)

	var reply cstructs.FsStatResponse
	var rpcErr error
	if localClient {
		rpcErr = s.agent.Client().ClientRPC("FileSystem.Stat", &args, &reply)
	} else if remoteClient {
		rpcErr = s.agent.Client().RPC("FileSystem.Stat", &args, &reply)
	} else if localServer {
		rpcErr = s.agent.Server().RPC("FileSystem.Stat", &args, &reply)
	}

	if rpcErr != nil {
		if structs.IsErrNoNodeConn(rpcErr) || structs.IsErrUnknownAllocation(rpcErr) {
			rpcErr = CodedError(404, rpcErr.Error())
		}

		return nil, rpcErr
	}

	return reply.Info, nil
}

func (s *HTTPServer) FileReadAtRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string
	var offset, limit int64
	var err error

	q := req.URL.Query()

	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/readat/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}
	if path = q.Get("path"); path == "" {
		return nil, fileNameNotPresentErr
	}

	if offset, err = strconv.ParseInt(q.Get("offset"), 10, 64); err != nil {
		return nil, fmt.Errorf("error parsing offset: %v", err)
	}

	// Parse the limit
	if limitStr := q.Get("limit"); limitStr != "" {
		if limit, err = strconv.ParseInt(limitStr, 10, 64); err != nil {
			return nil, fmt.Errorf("error parsing limit: %v", err)
		}
	}

	// Create the request arguments
	fsReq := &cstructs.FsStreamRequest{
		AllocID:   allocID,
		Path:      path,
		Offset:    offset,
		Origin:    "start",
		Limit:     limit,
		PlainText: true,
	}
	s.parse(resp, req, &fsReq.QueryOptions.Region, &fsReq.QueryOptions)

	// Make the request
	return s.fsStreamImpl(resp, req, "FileSystem.Stream", fsReq, fsReq.AllocID)
}

func (s *HTTPServer) FileCatRequest(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string

	q := req.URL.Query()

	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/cat/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}
	if path = q.Get("path"); path == "" {
		return nil, fileNameNotPresentErr
	}

	// Create the request arguments
	fsReq := &cstructs.FsStreamRequest{
		AllocID:   allocID,
		Path:      path,
		Origin:    "start",
		PlainText: true,
	}
	s.parse(resp, req, &fsReq.QueryOptions.Region, &fsReq.QueryOptions)

	// Make the request
	return s.fsStreamImpl(resp, req, "FileSystem.Stream", fsReq, fsReq.AllocID)
}

// Stream streams the content of a file blocking on EOF.
// The parameters are:
// * path: path to file to stream.
// * offset: The offset to start streaming data at, defaults to zero.
// * origin: Either "start" or "end" and defines from where the offset is
//           applied. Defaults to "start".
func (s *HTTPServer) Stream(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, path string

	q := req.URL.Query()

	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/stream/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}

	if path = q.Get("path"); path == "" {
		return nil, fileNameNotPresentErr
	}

	var offset int64
	offsetString := q.Get("offset")
	if offsetString != "" {
		var err error
		if offset, err = strconv.ParseInt(offsetString, 10, 64); err != nil {
			return nil, fmt.Errorf("error parsing offset: %v", err)
		}
	}

	origin := q.Get("origin")
	switch origin {
	case "start", "end":
	case "":
		origin = "start"
	default:
		return nil, invalidOrigin
	}

	// Create the request arguments
	fsReq := &cstructs.FsStreamRequest{
		AllocID: allocID,
		Path:    path,
		Origin:  origin,
		Offset:  offset,
		Follow:  true,
	}
	s.parse(resp, req, &fsReq.QueryOptions.Region, &fsReq.QueryOptions)

	// Make the request
	return s.fsStreamImpl(resp, req, "FileSystem.Stream", fsReq, fsReq.AllocID)
}

// Logs streams the content of a log blocking on EOF. The parameters are:
// * task: task name to stream logs for.
// * type: stdout/stderr to stream.
// * follow: A boolean of whether to follow the logs.
// * offset: The offset to start streaming data at, defaults to zero.
// * origin: Either "start" or "end" and defines from where the offset is
//           applied. Defaults to "start".
func (s *HTTPServer) Logs(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	var allocID, task, logType string
	var plain, follow bool
	var err error

	q := req.URL.Query()
	if allocID = strings.TrimPrefix(req.URL.Path, "/v1/client/fs/logs/"); allocID == "" {
		return nil, allocIDNotPresentErr
	}

	if task = q.Get("task"); task == "" {
		return nil, taskNotPresentErr
	}

	if followStr := q.Get("follow"); followStr != "" {
		if follow, err = strconv.ParseBool(followStr); err != nil {
			return nil, fmt.Errorf("Failed to parse follow field to boolean: %v", err)
		}
	}

	if plainStr := q.Get("plain"); plainStr != "" {
		if plain, err = strconv.ParseBool(plainStr); err != nil {
			return nil, fmt.Errorf("Failed to parse plain field to boolean: %v", err)
		}
	}

	logType = q.Get("type")
	switch logType {
	case "stdout", "stderr":
	default:
		return nil, logTypeNotPresentErr
	}

	var offset int64
	offsetString := q.Get("offset")
	if offsetString != "" {
		var err error
		if offset, err = strconv.ParseInt(offsetString, 10, 64); err != nil {
			return nil, fmt.Errorf("error parsing offset: %v", err)
		}
	}

	origin := q.Get("origin")
	switch origin {
	case "start", "end":
	case "":
		origin = "start"
	default:
		return nil, invalidOrigin
	}

	// Create the request arguments
	fsReq := &cstructs.FsLogsRequest{
		AllocID:   allocID,
		Task:      task,
		LogType:   logType,
		Offset:    offset,
		Origin:    origin,
		PlainText: plain,
		Follow:    follow,
	}
	s.parse(resp, req, &fsReq.QueryOptions.Region, &fsReq.QueryOptions)

	// Make the request
	return s.fsStreamImpl(resp, req, "FileSystem.Logs", fsReq, fsReq.AllocID)
}

// fsStreamImpl is used to make a streaming filesystem call that serializes the
// args and then expects a stream of StreamErrWrapper results where the payload
// is copied to the response body.
func (s *HTTPServer) fsStreamImpl(resp http.ResponseWriter,
	req *http.Request, method string, args interface{}, allocID string) (interface{}, error) {

	// Get the correct handler
	localClient, remoteClient, localServer := s.rpcHandlerForAlloc(allocID)
	var handler structs.StreamingRpcHandler
	var handlerErr error
	if localClient {
		handler, handlerErr = s.agent.Client().StreamingRpcHandler(method)
	} else if remoteClient {
		handler, handlerErr = s.agent.Client().RemoteStreamingRpcHandler(method)
	} else if localServer {
		handler, handlerErr = s.agent.Server().StreamingRpcHandler(method)
	}

	if handlerErr != nil {
		return nil, CodedError(500, handlerErr.Error())
	}

	// Create a pipe connecting the (possibly remote) handler to the http response
	httpPipe, handlerPipe := net.Pipe()
	decoder := codec.NewDecoder(httpPipe, structs.MsgpackHandle)
	encoder := codec.NewEncoder(httpPipe, structs.MsgpackHandle)

	// Create a goroutine that closes the pipe if the connection closes.
	ctx, cancel := context.WithCancel(req.Context())
	go func() {
		<-ctx.Done()
		httpPipe.Close()
	}()

	// Create an output that gets flushed on every write
	output := ioutils.NewWriteFlusher(resp)

	// Create a channel that decodes the results
	errCh := make(chan HTTPCodedError)
	go func() {
		defer cancel()

		// Send the request
		if err := encoder.Encode(args); err != nil {
			errCh <- CodedError(500, err.Error())
			return
		}
		encoder.Reset(httpPipe)

		for {
			select {
			case <-ctx.Done():
				errCh <- nil
				return
			default:
			}

			var res cstructs.StreamErrWrapper
			if err := decoder.Decode(&res); err != nil {
				errCh <- CodedError(500, err.Error())
				return
			}
			decoder.Reset(httpPipe)

			if err := res.Error; err != nil {
				if err.Code != nil {
					errCh <- CodedError(int(*err.Code), err.Error())
					return
				}
			}

			if _, err := io.Copy(output, bytes.NewReader(res.Payload)); err != nil {
				errCh <- CodedError(500, err.Error())
				return
			}
		}
	}()

	handler(handlerPipe)
	cancel()
	codedErr := <-errCh
	if codedErr != nil &&
		(codedErr == io.EOF ||
			strings.Contains(codedErr.Error(), "closed") ||
			strings.Contains(codedErr.Error(), "EOF")) {
		codedErr = nil
	}
	return nil, codedErr
}
<|endoftext|>


<commit_before>package wechat

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/KevinGong2013/ggbot/utils"
	"github.com/allegro/bigcache"
)

// ContactIterator iterator all contacts
type ContactIterator struct {
	it *bigcache.EntryInfoIterator
}

// Contact is wx Account struct
type Contact struct {
	UserName          string
	NickName          string
	HeadImgURL        string `json:"HeadImgUrl"`
	RemarkName        string
	PYInitial         string
	PYQuanPin         string
	RemarkPYInitial   string
	RemarkPYQuanPin   string
	HideInputBarFlag  float64
	StarFriend        float64
	Sex               float64
	Signature         string
	AppAccountFlag    float64
	VerifyFlag        float64
	ContactFlag       float64
	WebWxPluginSwitch float64
	HeadImgFlag       float64
	SnsFlag           float64
	Province          string
	City              string
	Alias             string
	DisplayName       string
	KeyWord           string
	EncryChatRoomID   string `json:"EncryChatRoomId"`
	IsOwner           float64
	Type              int
	ChangeType        int // 0 修改 1 删除
	MemberCount       float64
	MemberList        []*Contact
}

type updateGroupRequest struct {
	BaseRequest
	Count int
	List  []string
}

type updateGroupMemberRequest struct {
	BaseRequest
}

type getContactResponse struct {
	Response
	MemberCount int
	MemberList  []map[string]interface{}
	Seq         float64
}

type batchGetContactResponse struct {
	Response
	Count       int
	ContactList []map[string]interface{}
}

const (
	// ContactTypeFriend friend
	ContactTypeFriend = 1
	// ContactTypeGroup group
	ContactTypeGroup = 2
	// ContactTypeOfficial official
	ContactTypeOfficial = 3
)

var maxCountOnceLoadGroupMember = 50

// To is contact's ID can be used in msg struct
func (contact *Contact) To() string {
	return contact.UserName
}

func (wechat *WeChat) getContacts(seq float64) ([]map[string]interface{}, float64, error) {

	url := fmt.Sprintf(`%s/webwxgetcontact?%s&%s&r=%s&seq=%v`, wechat.BaseURL, wechat.PassTicketKV(), wechat.SkeyKV(), utils.Now(), seq)
	resp := new(getContactResponse)

	err := wechat.Excute(url, nil, resp)

	if err != nil {
		return nil, 0, err
	}

	return resp.MemberList, resp.Seq, nil
}

// SyncContact with Wechat server.
func (wechat *WeChat) SyncContact() error {
	wechat.resetCache()

	seq := float64(-1)

	for seq != 0 {
		if seq == -1 {
			seq = 0
		}
		memberList, s, err := wechat.getContacts(seq)
		if err != nil {
			return err
		}
		wechat.updateLocalContact(memberList)
		seq = s
	}

	return nil
}

func (wechat *WeChat) updateLocalContact(memberList []map[string]interface{}) {

	var gs []string

	for _, v := range memberList {

		vf, _ := v[`VerifyFlag`].(float64)
		un, _ := v[`UserName`].(string)

		if vf/8 != 0 {
			v[`Type`] = ContactTypeOfficial
		} else if strings.HasPrefix(un, `@@`) {
			v[`Type`] = ContactTypeGroup
			gs = append(gs, un)
		} else {
			v[`Type`] = ContactTypeFriend
		}

		wechat.saveContactToCache(v)
	}

	for _, g := range gs {
		wechat.FourceUpdateGroup(g)
	}
}

func (wechat *WeChat) fetchGroups(usernames []string) ([]map[string]interface{}, error) {

	var list []map[string]string
	for _, u := range usernames {
		list = append(list, map[string]string{
			`UserName`:   u,
			`ChatRoomId`: ``,
		})
	}

	data, err := json.Marshal(map[string]interface{}{
		`BaseRequest`: wechat.BaseRequest,
		`Count`:       len(list),
		`List`:        list,
	})
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf(`%s/webwxbatchgetcontact?type=ex&r=%v`, wechat.BaseURL, time.Now().Unix()*1000)
	resp := new(batchGetContactResponse)

	wechat.Excute(url, bytes.NewReader(data), resp)

	if !resp.IsSuccess() {
		logger.Error(resp.Error())
	}

	return resp.ContactList, nil
}

func (wechat *WeChat) fetchGroupsMembers(groups []map[string]interface{}) ([]map[string]interface{}, error) {

	list := make([]map[string]string, 0)

	for _, group := range groups {

		encryChatRoomID, _ := group[`EncryChatRoomId`].(string)
		members, _ := group[`MemberList`].([]interface{})

		logger.Debugf(`members %v`, members)

		for _, m := range members {
			mmap, _ := m.(map[string]interface{})
			u, _ := mmap[`UserName`].(string)
			list = append(list, map[string]string{
				`UserName`:        u,
				`EncryChatRoomId`: encryChatRoomID,
			})
		}
	}

	return wechat.fetchMembers(list), nil
}

func (wechat *WeChat) fetchMembers(list []map[string]string) []map[string]interface{} {

	if len(list) > maxCountOnceLoadGroupMember {
		return append(wechat.fetchMembers(list[:maxCountOnceLoadGroupMember]), wechat.fetchMembers(list[maxCountOnceLoadGroupMember:len(list)])...)
	}

	data, _ := json.Marshal(map[string]interface{}{
		`BaseRequest`: wechat.BaseRequest,
		`Count`:       len(list),
		`List`:        list,
	})

	url := fmt.Sprintf(`%s/webwxbatchgetcontact?type=ex&r=%v`, wechat.BaseURL, time.Now().Unix()*1000)
	resp := new(batchGetContactResponse)

	wechat.Excute(url, bytes.NewReader(data), resp)

	if !resp.IsSuccess() {
		logger.Error(resp.Error())
	}

	return resp.ContactList
}

// UpateGroupIfNeeded ...
func (wechat *WeChat) UpateGroupIfNeeded(groupID string) {

	bs, _ := wechat.contactCache.Get(groupID)

	if bs == nil {
		wechat.FourceUpdateGroup(groupID)
	}
}

// FourceUpdateGroup upate group infomation
func (wechat *WeChat) FourceUpdateGroup(groupID string) {

	groups, err := wechat.fetchGroups([]string{groupID})
	if err != nil || len(groups) != 1 {
		logger.Error(`sync group failed`)
		return
	}

	// 保存群组
	for _, v := range groups {
		v[`Type`] = ContactTypeGroup
		wechat.saveContactToCache(v)
	}

	memberList, err := wechat.fetchGroupsMembers(groups)
	if err != nil {
		logger.Error(`sync group failed`)
		return
	}

	for _, v := range memberList {
		v[`Type`] = ContactTypeFriend
		wechat.saveContactToCache(v)
	}
}

// ContactByUserName ...
func (wechat *WeChat) ContactByUserName(un string) (*Contact, error) {
	bs, err := wechat.contactCache.Get(un)
	if err != nil {
		return nil, err
	}
	var contact *Contact
	err = json.NewDecoder(bytes.NewReader(bs)).Decode(&contact)
	if err != nil {
		return nil, err
	}
	return contact, nil
}

// UserNameByNickName ..
func (wechat *WeChat) UserNameByNickName(nn string) (string, error) {

	bs, err := wechat.nicknameCache.Get(nn)
	if err != nil {
		return ``, nil
	}
	return string(bs), nil
}

// ContactByNickName search contact with nick name
func (wechat *WeChat) ContactByNickName(nn string) (*Contact, error) {
	un, err := wechat.UserNameByNickName(nn)
	if err != nil {
		return nil, err
	}
	return wechat.ContactByUserName(un)
}

// ContactIterator use this iterator all contacts // TODO
func (wechat *WeChat) ContactIterator() *ContactIterator {
	return &ContactIterator{wechat.contactCache.Iterator()}
}

// TODO
func (wechat *WeChat) modifyRemarkName(un string) (string, error) {

	data, _ := json.Marshal(map[string]interface{}{
		`BaseRequest`: wechat.BaseRequest,
		`UserName`:    un,
		`CmdId`:       2,
		`NickName`:    `Test`,
	})

	url := fmt.Sprintf(`%s/webwxoplog?lang=zh_CN&%v`, wechat.BaseURL, wechat.PassTicketKV())
	resp := new(Response)

	wechat.Excute(url, bytes.NewReader(data), resp)

	if !resp.IsSuccess() {
		logger.Error(resp.Error())
	}

	return `Test`, nil
}

func (wechat *WeChat) saveContactToCache(contact map[string]interface{}) {
	un, _ := contact[`UserName`].(string)
	bs, e := json.Marshal(contact)
	if e != nil {
		logger.Error(e)
	} else {
		wechat.contactCache.Set(un, bs)
		nn, _ := contact[`NickName`].(string)
		wechat.nicknameCache.Set(nn, []byte(un))
	}
}

func (wechat *WeChat) resetCache() {

	wechat.contactCache.Reset()
	wechat.nicknameCache.Reset()
}

func (wechat *WeChat) contactDidChange(cts *CountedContent, changeType int) {
	if changeType == 0 { // 修改
		for _, v := range cts.Content {
			vf, _ := v[`VerifyFlag`].(float64)
			un, _ := v[`UserName`].(string)

			if vf/8 != 0 {
				v[`Type`] = ContactTypeOfficial
			} else if strings.HasPrefix(un, `@@`) {
				v[`Type`] = ContactTypeGroup
			} else {
				v[`Type`] = ContactTypeFriend
			}

			wechat.saveContactToCache(v)
		}
	}
}
<commit_msg>Refine code.<commit_after>package wechat

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/KevinGong2013/ggbot/utils"
	"github.com/allegro/bigcache"
)

// ContactIterator iterator all contacts
type ContactIterator struct {
	it *bigcache.EntryInfoIterator
}

// Contact is wx Account struct
type Contact struct {
	UserName          string
	NickName          string
	HeadImgURL        string `json:"HeadImgUrl"`
	RemarkName        string
	PYInitial         string
	PYQuanPin         string
	RemarkPYInitial   string
	RemarkPYQuanPin   string
	HideInputBarFlag  float64
	StarFriend        float64
	Sex               float64
	Signature         string
	AppAccountFlag    float64
	VerifyFlag        float64
	ContactFlag       float64
	WebWxPluginSwitch float64
	HeadImgFlag       float64
	SnsFlag           float64
	Province          string
	City              string
	Alias             string
	DisplayName       string
	KeyWord           string
	EncryChatRoomID   string `json:"EncryChatRoomId"`
	IsOwner           float64
	Type              int
	ChangeType        int // 0 修改 1 删除
	MemberCount       float64
	MemberList        []*Contact
}

type updateGroupRequest struct {
	BaseRequest
	Count int
	List  []string
}

type updateGroupMemberRequest struct {
	BaseRequest
}

type getContactResponse struct {
	Response
	MemberCount int
	MemberList  []map[string]interface{}
	Seq         float64
}

type batchGetContactResponse struct {
	Response
	Count       int
	ContactList []map[string]interface{}
}

const (
	// ContactTypeFriend friend
	ContactTypeFriend = 1
	// ContactTypeGroup group
	ContactTypeGroup = 2
	// ContactTypeOfficial official
	ContactTypeOfficial = 3
)

var maxCountOnceLoadGroupMember = 50

// To is contact's ID can be used in msg struct
func (contact *Contact) To() string {
	return contact.UserName
}

func (wechat *WeChat) getContacts(seq float64) ([]map[string]interface{}, float64, error) {

	url := fmt.Sprintf(`%s/webwxgetcontact?%s&%s&r=%s&seq=%v`, wechat.BaseURL, wechat.PassTicketKV(), wechat.SkeyKV(), utils.Now(), seq)
	resp := new(getContactResponse)

	err := wechat.Excute(url, nil, resp)

	if err != nil {
		return nil, 0, err
	}

	return resp.MemberList, resp.Seq, nil
}

// SyncContact with Wechat server.
func (wechat *WeChat) SyncContact() error {
	wechat.resetCache()

	seq := float64(-1)

	for seq != 0 {
		if seq == -1 {
			seq = 0
		}
		memberList, s, err := wechat.getContacts(seq)
		if err != nil {
			return err
		}
		wechat.updateLocalContact(memberList)
		seq = s
	}

	return nil
}

func (wechat *WeChat) updateLocalContact(memberList []map[string]interface{}) {

	var gs []string

	for _, v := range memberList {

		vf, _ := v[`VerifyFlag`].(float64)
		un, _ := v[`UserName`].(string)

		if vf/8 != 0 {
			v[`Type`] = ContactTypeOfficial
		} else if strings.HasPrefix(un, `@@`) {
			v[`Type`] = ContactTypeGroup
			gs = append(gs, un)
		} else {
			v[`Type`] = ContactTypeFriend
		}

		wechat.saveContactToCache(v)
	}

	for _, g := range gs {
		wechat.FourceUpdateGroup(g)
	}
}

func (wechat *WeChat) fetchGroups(usernames []string) ([]map[string]interface{}, error) {

	var list []map[string]string
	for _, u := range usernames {
		list = append(list, map[string]string{
			`UserName`:   u,
			`ChatRoomId`: ``,
		})
	}

	data, err := json.Marshal(map[string]interface{}{
		`BaseRequest`: wechat.BaseRequest,
		`Count`:       len(list),
		`List`:        list,
	})
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf(`%s/webwxbatchgetcontact?type=ex&r=%v`, wechat.BaseURL, time.Now().Unix()*1000)
	resp := new(batchGetContactResponse)

	wechat.Excute(url, bytes.NewReader(data), resp)

	if !resp.IsSuccess() {
		logger.Error(resp.Error())
	}

	return resp.ContactList, nil
}

func (wechat *WeChat) fetchGroupsMembers(groups []map[string]interface{}) ([]map[string]interface{}, error) {

	list := make([]map[string]string, 0)

	for _, group := range groups {

		encryChatRoomID, _ := group[`EncryChatRoomId`].(string)
		members, _ := group[`MemberList`].([]interface{})

		for _, m := range members {
			mmap, _ := m.(map[string]interface{})
			u, _ := mmap[`UserName`].(string)
			list = append(list, map[string]string{
				`UserName`:        u,
				`EncryChatRoomId`: encryChatRoomID,
			})
		}
	}

	return wechat.fetchMembers(list), nil
}

func (wechat *WeChat) fetchMembers(list []map[string]string) []map[string]interface{} {

	if len(list) > maxCountOnceLoadGroupMember {
		return append(wechat.fetchMembers(list[:maxCountOnceLoadGroupMember]), wechat.fetchMembers(list[maxCountOnceLoadGroupMember:len(list)])...)
	}

	data, _ := json.Marshal(map[string]interface{}{
		`BaseRequest`: wechat.BaseRequest,
		`Count`:       len(list),
		`List`:        list,
	})

	url := fmt.Sprintf(`%s/webwxbatchgetcontact?type=ex&r=%v`, wechat.BaseURL, time.Now().Unix()*1000)
	resp := new(batchGetContactResponse)

	wechat.Excute(url, bytes.NewReader(data), resp)

	if !resp.IsSuccess() {
		logger.Error(resp.Error())
	}

	return resp.ContactList
}

// UpateGroupIfNeeded ...
func (wechat *WeChat) UpateGroupIfNeeded(groupID string) {

	bs, _ := wechat.contactCache.Get(groupID)

	if bs == nil {
		wechat.FourceUpdateGroup(groupID)
	}
}

// FourceUpdateGroup upate group infomation
func (wechat *WeChat) FourceUpdateGroup(groupID string) {

	groups, err := wechat.fetchGroups([]string{groupID})
	if err != nil || len(groups) != 1 {
		logger.Error(`sync group failed`)
		return
	}

	// 保存群组
	for _, v := range groups {
		v[`Type`] = ContactTypeGroup
		wechat.saveContactToCache(v)
	}

	memberList, err := wechat.fetchGroupsMembers(groups)
	if err != nil {
		logger.Error(`sync group failed`)
		return
	}

	for _, v := range memberList {
		v[`Type`] = ContactTypeFriend
		wechat.saveContactToCache(v)
	}
}

// ContactByUserName ...
func (wechat *WeChat) ContactByUserName(un string) (*Contact, error) {
	bs, err := wechat.contactCache.Get(un)
	if err != nil {
		return nil, err
	}
	var contact *Contact
	err = json.NewDecoder(bytes.NewReader(bs)).Decode(&contact)
	if err != nil {
		return nil, err
	}
	return contact, nil
}

// UserNameByNickName ..
func (wechat *WeChat) UserNameByNickName(nn string) (string, error) {

	bs, err := wechat.nicknameCache.Get(nn)
	if err != nil {
		return ``, nil
	}
	return string(bs), nil
}

// ContactByNickName search contact with nick name
func (wechat *WeChat) ContactByNickName(nn string) (*Contact, error) {
	un, err := wechat.UserNameByNickName(nn)
	if err != nil {
		return nil, err
	}
	return wechat.ContactByUserName(un)
}

// ContactIterator use this iterator all contacts // TODO
func (wechat *WeChat) ContactIterator() *ContactIterator {
	return &ContactIterator{wechat.contactCache.Iterator()}
}

// TODO
func (wechat *WeChat) modifyRemarkName(un string) (string, error) {

	data, _ := json.Marshal(map[string]interface{}{
		`BaseRequest`: wechat.BaseRequest,
		`UserName`:    un,
		`CmdId`:       2,
		`NickName`:    `Test`,
	})

	url := fmt.Sprintf(`%s/webwxoplog?lang=zh_CN&%v`, wechat.BaseURL, wechat.PassTicketKV())
	resp := new(Response)

	wechat.Excute(url, bytes.NewReader(data), resp)

	if !resp.IsSuccess() {
		logger.Error(resp.Error())
	}

	return `Test`, nil
}

func (wechat *WeChat) saveContactToCache(contact map[string]interface{}) {
	un, _ := contact[`UserName`].(string)
	bs, e := json.Marshal(contact)
	if e != nil {
		logger.Error(e)
	} else {
		wechat.contactCache.Set(un, bs)
		nn, _ := contact[`NickName`].(string)
		wechat.nicknameCache.Set(nn, []byte(un))
	}
}

func (wechat *WeChat) resetCache() {

	wechat.contactCache.Reset()
	wechat.nicknameCache.Reset()
}

func (wechat *WeChat) contactDidChange(cts *CountedContent, changeType int) {
	if changeType == 0 { // 修改
		for _, v := range cts.Content {
			vf, _ := v[`VerifyFlag`].(float64)
			un, _ := v[`UserName`].(string)

			if vf/8 != 0 {
				v[`Type`] = ContactTypeOfficial
			} else if strings.HasPrefix(un, `@@`) {
				v[`Type`] = ContactTypeGroup
			} else {
				v[`Type`] = ContactTypeFriend
			}

			wechat.saveContactToCache(v)
		}
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/bleenco/abstruse/internal/auth"
	"github.com/bleenco/abstruse/internal/version"
	"github.com/bleenco/abstruse/pkg/fs"
	"github.com/bleenco/abstruse/pkg/lib"
	"github.com/bleenco/abstruse/pkg/tlsutil"
	"github.com/bleenco/abstruse/worker/app"
	"github.com/bleenco/abstruse/worker/config"
	"github.com/bleenco/abstruse/worker/docker"
	"github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"go.uber.org/zap"
)

var (
	cfgFile string
	rootCmd = &cobra.Command{
		Use:   "abstruse-worker",
		Short: "Abstruse CI Worker Node",
		Run: func(cmd *cobra.Command, args []string) {
			app, err := CreateApp()
			if err != nil {
				fatal(err)
			}

			if err := app.run(); err != nil {
				fatal(err)
			}
			os.Exit(0)
		},
	}
	versionCmd = &cobra.Command{
		Use:   "version",
		Short: "Print the version number and build info",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println(version.GenerateBuildVersionString())
			os.Exit(0)
		},
	}
)

type application struct {
	config *config.Config
	logger *zap.Logger
	app    *app.App
}

func newApplication(
	config *config.Config,
	logger *zap.Logger,
	app *app.App,
) *application {
	return &application{config, logger, app}
}

func (a application) run() error {
	errch := make(chan error, 1)

	go func() {
		if err := a.app.Run(); err != nil {
			errch <- err
		}
	}()

	return <-errch
}

// Execute executes the root command.
func Execute() error {
	return rootCmd.Execute()
}

func init() {
	rootCmd.AddCommand(versionCmd)
	cobra.OnInitialize(initDefaults)

	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/abstruse/abstruse-worker.json)")
	rootCmd.PersistentFlags().String("id", lib.RandomString(), "worker node ID")
	rootCmd.PersistentFlags().String("server-addr", "0.0.0.0:6500", "abstruse server remote address")
	rootCmd.PersistentFlags().String("grpc-addr", "0.0.0.0:3330", "gRPC server listen address")
	rootCmd.PersistentFlags().String("tls-cert", "cert-worker.pem", "path to SSL certificate file")
	rootCmd.PersistentFlags().String("tls-key", "key-worker.pem", "path to SSL private key file")
	rootCmd.PersistentFlags().Int("scheduler-maxparallel", 5, "scheduler max parallel option defines how many jobs can run in parallel")
	rootCmd.PersistentFlags().String("auth-jwtsecret", lib.RandomString(), "JWT authentication secret key")
	rootCmd.PersistentFlags().String("registry-addr", "https://registry-1.docker.io", "docker image registry server addr")
	rootCmd.PersistentFlags().String("registry-username", "", "docker image registry username")
	rootCmd.PersistentFlags().String("registry-password", "", "docker image registry password")
	rootCmd.PersistentFlags().String("logger-level", "info", "logging level (available options: debug, info, warn, error, panic, fatal)")
	rootCmd.PersistentFlags().Bool("logger-stdout", true, "print logs to stdout")
	rootCmd.PersistentFlags().String("logger-filename", "abstruse-worker.log", "log filename")
	rootCmd.PersistentFlags().Int("logger-max-size", 500, "maximum log file size (in MB)")
	rootCmd.PersistentFlags().Int("logger-max-backups", 3, "maximum log file backups")
	rootCmd.PersistentFlags().Int("logger-max-age", 3, "maximum log age")
}

func initDefaults() {
	viper.BindPFlag("grpc.addr", rootCmd.PersistentFlags().Lookup("grpc-addr"))
	viper.BindPFlag("id", rootCmd.PersistentFlags().Lookup("id"))
	viper.BindPFlag("server.addr", rootCmd.PersistentFlags().Lookup("server-addr"))
	viper.BindPFlag("tls.cert", rootCmd.PersistentFlags().Lookup("tls-cert"))
	viper.BindPFlag("tls.key", rootCmd.PersistentFlags().Lookup("tls-key"))
	viper.BindPFlag("scheduler.maxparallel", rootCmd.PersistentFlags().Lookup("scheduler-maxparallel"))
	viper.BindPFlag("auth.jwtsecret", rootCmd.PersistentFlags().Lookup("auth-jwtsecret"))
	viper.BindPFlag("registry.addr", rootCmd.PersistentFlags().Lookup("registry-addr"))
	viper.BindPFlag("registry.username", rootCmd.PersistentFlags().Lookup("registry-username"))
	viper.BindPFlag("registry.password", rootCmd.PersistentFlags().Lookup("registry-password"))
	viper.BindPFlag("logger.level", rootCmd.PersistentFlags().Lookup("logger-level"))
	viper.BindPFlag("logger.stdout", rootCmd.PersistentFlags().Lookup("logger-stdout"))
	viper.BindPFlag("logger.filename", rootCmd.PersistentFlags().Lookup("logger-filename"))
	viper.BindPFlag("logger.maxsize", rootCmd.PersistentFlags().Lookup("logger-max-size"))
	viper.BindPFlag("logger.maxbackups", rootCmd.PersistentFlags().Lookup("logger-max-backups"))
	viper.BindPFlag("logger.maxage", rootCmd.PersistentFlags().Lookup("logger-max-age"))
}

func newConfig() *config.Config {
	var cfg *config.Config

	if cfgFile == "" {
		home, err := homedir.Dir()
		if err != nil {
			fatal(err)
		}
		cfgFile = filepath.Join(home, "abstruse", "abstruse-worker.json")
	}

	viper.SetConfigFile(cfgFile)
	viper.SetConfigType("json")
	viper.SetEnvPrefix("abstruse")
	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
	viper.AutomaticEnv()

	cfgFileUsed := viper.ConfigFileUsed()

	if !fs.Exists(cfgFileUsed) {
		if !fs.Exists(filepath.Dir(cfgFileUsed)) {
			if err := fs.MakeDir(filepath.Dir(cfgFileUsed)); err != nil {
				fatal(err)
			}
		}

		if err := viper.SafeWriteConfigAs(cfgFileUsed); err != nil {
			fatal(err)
		}
	}

	if err := viper.ReadInConfig(); err != nil {
		fatal(err)
	}

	if err := viper.Unmarshal(&cfg); err != nil {
		fatal(err)
	}

	if !strings.HasPrefix(cfg.Logger.Filename, "/") {
		cfg.Logger.Filename = filepath.Join(filepath.Dir(viper.ConfigFileUsed()), cfg.Logger.Filename)
	}

	if !strings.HasPrefix(cfg.TLS.Cert, "/") {
		cfg.TLS.Cert = filepath.Join(filepath.Dir(cfgFileUsed), cfg.TLS.Cert)
	}

	if !strings.HasPrefix(cfg.TLS.Key, "/") {
		cfg.TLS.Key = filepath.Join(filepath.Dir(cfgFileUsed), cfg.TLS.Key)
	}

	auth.Init(viper.GetString("auth.jwtsecret"))

	cert, key := cfg.TLS.Cert, cfg.TLS.Key
	if !strings.HasPrefix(cert, "/") {
		cert = filepath.Join(filepath.Dir(viper.ConfigFileUsed()), cert)
	}
	if !strings.HasPrefix(key, "/") {
		key = filepath.Join(filepath.Dir(viper.ConfigFileUsed()), key)
	}

	if err := tlsutil.CheckAndGenerateCert(cert, key); err != nil {
		fatal(err)
	}

	docker.Init(cfg.Registry)

	return cfg
}

func fatal(msg interface{}) {
	fmt.Println(msg)
	os.Exit(1)
}
<commit_msg>chore(worker): update default server api addr<commit_after>package cmd

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/bleenco/abstruse/internal/auth"
	"github.com/bleenco/abstruse/internal/version"
	"github.com/bleenco/abstruse/pkg/fs"
	"github.com/bleenco/abstruse/pkg/lib"
	"github.com/bleenco/abstruse/pkg/tlsutil"
	"github.com/bleenco/abstruse/worker/app"
	"github.com/bleenco/abstruse/worker/config"
	"github.com/bleenco/abstruse/worker/docker"
	"github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"go.uber.org/zap"
)

var (
	cfgFile string
	rootCmd = &cobra.Command{
		Use:   "abstruse-worker",
		Short: "Abstruse CI Worker Node",
		Run: func(cmd *cobra.Command, args []string) {
			app, err := CreateApp()
			if err != nil {
				fatal(err)
			}

			if err := app.run(); err != nil {
				fatal(err)
			}
			os.Exit(0)
		},
	}
	versionCmd = &cobra.Command{
		Use:   "version",
		Short: "Print the version number and build info",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println(version.GenerateBuildVersionString())
			os.Exit(0)
		},
	}
)

type application struct {
	config *config.Config
	logger *zap.Logger
	app    *app.App
}

func newApplication(
	config *config.Config,
	logger *zap.Logger,
	app *app.App,
) *application {
	return &application{config, logger, app}
}

func (a application) run() error {
	errch := make(chan error, 1)

	go func() {
		if err := a.app.Run(); err != nil {
			errch <- err
		}
	}()

	return <-errch
}

// Execute executes the root command.
func Execute() error {
	return rootCmd.Execute()
}

func init() {
	rootCmd.AddCommand(versionCmd)
	cobra.OnInitialize(initDefaults)

	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/abstruse/abstruse-worker.json)")
	rootCmd.PersistentFlags().String("id", lib.RandomString(), "worker node ID")
	rootCmd.PersistentFlags().String("server-addr", "http://localhost", "abstruse server API address")
	rootCmd.PersistentFlags().String("grpc-addr", "0.0.0.0:3330", "gRPC server listen address")
	rootCmd.PersistentFlags().String("tls-cert", "cert-worker.pem", "path to SSL certificate file")
	rootCmd.PersistentFlags().String("tls-key", "key-worker.pem", "path to SSL private key file")
	rootCmd.PersistentFlags().Int("scheduler-maxparallel", 5, "scheduler max parallel option defines how many jobs can run in parallel")
	rootCmd.PersistentFlags().String("auth-jwtsecret", lib.RandomString(), "JWT authentication secret key")
	rootCmd.PersistentFlags().String("registry-addr", "https://registry-1.docker.io", "docker image registry server addr")
	rootCmd.PersistentFlags().String("registry-username", "", "docker image registry username")
	rootCmd.PersistentFlags().String("registry-password", "", "docker image registry password")
	rootCmd.PersistentFlags().String("logger-level", "info", "logging level (available options: debug, info, warn, error, panic, fatal)")
	rootCmd.PersistentFlags().Bool("logger-stdout", true, "print logs to stdout")
	rootCmd.PersistentFlags().String("logger-filename", "abstruse-worker.log", "log filename")
	rootCmd.PersistentFlags().Int("logger-max-size", 500, "maximum log file size (in MB)")
	rootCmd.PersistentFlags().Int("logger-max-backups", 3, "maximum log file backups")
	rootCmd.PersistentFlags().Int("logger-max-age", 3, "maximum log age")
}

func initDefaults() {
	viper.BindPFlag("grpc.addr", rootCmd.PersistentFlags().Lookup("grpc-addr"))
	viper.BindPFlag("id", rootCmd.PersistentFlags().Lookup("id"))
	viper.BindPFlag("server.addr", rootCmd.PersistentFlags().Lookup("server-addr"))
	viper.BindPFlag("tls.cert", rootCmd.PersistentFlags().Lookup("tls-cert"))
	viper.BindPFlag("tls.key", rootCmd.PersistentFlags().Lookup("tls-key"))
	viper.BindPFlag("scheduler.maxparallel", rootCmd.PersistentFlags().Lookup("scheduler-maxparallel"))
	viper.BindPFlag("auth.jwtsecret", rootCmd.PersistentFlags().Lookup("auth-jwtsecret"))
	viper.BindPFlag("registry.addr", rootCmd.PersistentFlags().Lookup("registry-addr"))
	viper.BindPFlag("registry.username", rootCmd.PersistentFlags().Lookup("registry-username"))
	viper.BindPFlag("registry.password", rootCmd.PersistentFlags().Lookup("registry-password"))
	viper.BindPFlag("logger.level", rootCmd.PersistentFlags().Lookup("logger-level"))
	viper.BindPFlag("logger.stdout", rootCmd.PersistentFlags().Lookup("logger-stdout"))
	viper.BindPFlag("logger.filename", rootCmd.PersistentFlags().Lookup("logger-filename"))
	viper.BindPFlag("logger.maxsize", rootCmd.PersistentFlags().Lookup("logger-max-size"))
	viper.BindPFlag("logger.maxbackups", rootCmd.PersistentFlags().Lookup("logger-max-backups"))
	viper.BindPFlag("logger.maxage", rootCmd.PersistentFlags().Lookup("logger-max-age"))
}

func newConfig() *config.Config {
	var cfg *config.Config

	if cfgFile == "" {
		home, err := homedir.Dir()
		if err != nil {
			fatal(err)
		}
		cfgFile = filepath.Join(home, "abstruse", "abstruse-worker.json")
	}

	viper.SetConfigFile(cfgFile)
	viper.SetConfigType("json")
	viper.SetEnvPrefix("abstruse")
	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
	viper.AutomaticEnv()

	cfgFileUsed := viper.ConfigFileUsed()

	if !fs.Exists(cfgFileUsed) {
		if !fs.Exists(filepath.Dir(cfgFileUsed)) {
			if err := fs.MakeDir(filepath.Dir(cfgFileUsed)); err != nil {
				fatal(err)
			}
		}

		if err := viper.SafeWriteConfigAs(cfgFileUsed); err != nil {
			fatal(err)
		}
	}

	if err := viper.ReadInConfig(); err != nil {
		fatal(err)
	}

	if err := viper.Unmarshal(&cfg); err != nil {
		fatal(err)
	}

	if !strings.HasPrefix(cfg.Logger.Filename, "/") {
		cfg.Logger.Filename = filepath.Join(filepath.Dir(viper.ConfigFileUsed()), cfg.Logger.Filename)
	}

	if !strings.HasPrefix(cfg.TLS.Cert, "/") {
		cfg.TLS.Cert = filepath.Join(filepath.Dir(cfgFileUsed), cfg.TLS.Cert)
	}

	if !strings.HasPrefix(cfg.TLS.Key, "/") {
		cfg.TLS.Key = filepath.Join(filepath.Dir(cfgFileUsed), cfg.TLS.Key)
	}

	auth.Init(viper.GetString("auth.jwtsecret"))

	cert, key := cfg.TLS.Cert, cfg.TLS.Key
	if !strings.HasPrefix(cert, "/") {
		cert = filepath.Join(filepath.Dir(viper.ConfigFileUsed()), cert)
	}
	if !strings.HasPrefix(key, "/") {
		key = filepath.Join(filepath.Dir(viper.ConfigFileUsed()), key)
	}

	if err := tlsutil.CheckAndGenerateCert(cert, key); err != nil {
		fatal(err)
	}

	docker.Init(cfg.Registry)

	return cfg
}

func fatal(msg interface{}) {
	fmt.Println(msg)
	os.Exit(1)
}
<|endoftext|>


<commit_before>package workspace

import (
	"fmt"
	"strings"
	"time"
)

// A Priority allows tasks to be prioritised by their importance.
type Priority uint8

const (
	// PriorityUnknown is an invalid priority.
	PriorityUnknown Priority = iota

	// PriorityLow is intended for "rainy-day" tasks.
	PriorityLow

	// PriorityNormal is the default priority.
	PriorityNormal

	// PriorityHigh is intended for tasks that should be done
	// before normal-priority tasks.
	PriorityHigh

	// PriorityUrgent is intended for time-sensitive tasks.
	PriorityUrgent
)

var priorityStrings = map[Priority]string{
	PriorityUnknown: "?",
	PriorityLow:     "L",
	PriorityNormal:  "N",
	PriorityHigh:    "H",
	PriorityUrgent:  "!",
}

// String provides a string representation for the Priority type.
func (pri Priority) String() string {
	s, ok := priorityStrings[pri]
	if !ok {
		s = priorityStrings[PriorityUnknown]
	}
	return s
}

// PriorityFromString returns the appropriate Priority from a string.
func PriorityFromString(ps string) Priority {
	var pri Priority
	var s string

	for pri, s = range priorityStrings {
		if s == ps {
			return pri
		}
	}

	return PriorityUnknown
}

// PriorityStrings is a list of priority strings and their values,
// useful for usage messages.
var PriorityStrings = `Priority specifiers:

        ?       Unknown
        L       Low
        N       Normal
        H       High
        !       Urgent
`

// A Task is a TODO item.
type Task struct {
	ID                uint64
	Done              bool
	Created, Finished time.Time
	Title             string
	Notes             []string
	Tags              []string
	Priority          Priority
}

// String provides a default representation for a task.
func (t *Task) String() string {
	marker := " "
	if t.Done {
		marker = "X"
	}

	endDate := ""
	if t.Done {
		endDate = fmt.Sprintf(", completed %s", t.Finished.Format(DateFormat))
	}

	return fmt.Sprintf("[%s] %s (%s) - %s%s", marker, t.Title, t.Priority,
		t.Created.Format(DateFormat), endDate)
}

// NewTask returns a new incomplete task started now.
func NewTask(id uint64, title string) *Task {
	return &Task{
		ID:       id,
		Created:  time.Now(),
		Title:    title,
		Priority: PriorityNormal,
	}
}

// MarkDone marks a task as completed, marking the completion time as
// now.
func (t *Task) MarkDone() {
	t.Done = true
	t.Finished = time.Now()
}

// TagString returns a string containing all the tags in the task.
func (t *Task) TagString() string {
	return strings.Join(t.Tags, ", ")
}

// A TaskSet contains a set of tasks.
type TaskSet map[uint64]*Task

// FilterPriority returns all the tasks with at least the given priority.
func (ts TaskSet) FilterPriority(pri Priority) TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if task.Priority >= pri {
			tasks[id] = task
		}
	}

	return tasks
}

// FilterTag returns all tasks with the given tag.
func (ts TaskSet) FilterTag(tag string) TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if contains(tag, task.Tags) {
			tasks[id] = task
		}
	}

	return tasks
}

// Unfinished returns the subset of tasks that aren't completed.
func (ts TaskSet) Unfinished() TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if !task.Done {
			tasks[id] = task
		}
	}

	return tasks
}

// CompletedDuration returns the tasks completed within the last
// duration, selecting on the completion date.
func (ts TaskSet) CompletedDuration(dur time.Duration) TaskSet {
	var tasks = TaskSet{}

	started := time.Now().Add(-1 * dur)
	for id, task := range ts {
		if task.Done {
			if task.Finished.After(started) {
				tasks[id] = task
			}
		}
	}

	return tasks
}

// CreatedDuration returns the tasks completed within the last
// duration, selecting on the creation date.
func (ts TaskSet) CreatedDuration(dur time.Duration) TaskSet {
	var tasks = TaskSet{}

	started := time.Now().Add(-1 * dur)
	for id, task := range ts {
		if task.Done {
			if task.Created.After(started) {
				tasks[id] = task
			}
		}
	}

	return tasks
}

// CompletedRange returns a list of tasks completed within the
// specified times, selecting on the completion date.
func (ts TaskSet) CompletedRange(start, end time.Time) TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if task.Done {
			if after(task.Finished, start) && before(task.Finished, end) {
				tasks[id] = task
			}
		}
	}

	return tasks
}

// CreatedRange returns a list of tasks completed within the
// specified times, selecting on the created date.
func (ts TaskSet) CreatedRange(start, end time.Time) TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if task.Done {
			if task.Created.After(start) && task.Created.Before(end) {
				tasks[id] = task
			}
		}
	}

	return tasks
}

// Sort returns a list of the tasks in chronological order.
func (ts TaskSet) Sort() []*Task {
	var ids = make([]uint64, 0, len(ts))
	for id := range ts {
		ids = append(ids, id)
	}

	var tasks = make([]*Task, len(ids))
	for i := range ids {
		tasks[i] = ts[ids[i]]
	}

	return tasks
}

// NewTaskID returns a new task identifier.
func NewTaskID() uint64 {
	return uint64(time.Now().UnixNano())
}
<commit_msg>Add FilterTags method to TaskSet.<commit_after>package workspace

import (
	"fmt"
	"strings"
	"time"
)

// A Priority allows tasks to be prioritised by their importance.
type Priority uint8

const (
	// PriorityUnknown is an invalid priority.
	PriorityUnknown Priority = iota

	// PriorityLow is intended for "rainy-day" tasks.
	PriorityLow

	// PriorityNormal is the default priority.
	PriorityNormal

	// PriorityHigh is intended for tasks that should be done
	// before normal-priority tasks.
	PriorityHigh

	// PriorityUrgent is intended for time-sensitive tasks.
	PriorityUrgent
)

var priorityStrings = map[Priority]string{
	PriorityUnknown: "?",
	PriorityLow:     "L",
	PriorityNormal:  "N",
	PriorityHigh:    "H",
	PriorityUrgent:  "!",
}

// String provides a string representation for the Priority type.
func (pri Priority) String() string {
	s, ok := priorityStrings[pri]
	if !ok {
		s = priorityStrings[PriorityUnknown]
	}
	return s
}

// PriorityFromString returns the appropriate Priority from a string.
func PriorityFromString(ps string) Priority {
	var pri Priority
	var s string

	for pri, s = range priorityStrings {
		if s == ps {
			return pri
		}
	}

	return PriorityUnknown
}

// PriorityStrings is a list of priority strings and their values,
// useful for usage messages.
var PriorityStrings = `Priority specifiers:

        ?       Unknown
        L       Low
        N       Normal
        H       High
        !       Urgent
`

// A Task is a TODO item.
type Task struct {
	ID                uint64
	Done              bool
	Created, Finished time.Time
	Title             string
	Notes             []string
	Tags              []string
	Priority          Priority
}

// String provides a default representation for a task.
func (t *Task) String() string {
	marker := " "
	if t.Done {
		marker = "X"
	}

	endDate := ""
	if t.Done {
		endDate = fmt.Sprintf(", completed %s", t.Finished.Format(DateFormat))
	}

	return fmt.Sprintf("[%s] %s (%s) - %s%s", marker, t.Title, t.Priority,
		t.Created.Format(DateFormat), endDate)
}

// NewTask returns a new incomplete task started now.
func NewTask(id uint64, title string) *Task {
	return &Task{
		ID:       id,
		Created:  time.Now(),
		Title:    title,
		Priority: PriorityNormal,
	}
}

// MarkDone marks a task as completed, marking the completion time as
// now.
func (t *Task) MarkDone() {
	t.Done = true
	t.Finished = time.Now()
}

// TagString returns a string containing all the tags in the task.
func (t *Task) TagString() string {
	return strings.Join(t.Tags, ", ")
}

// A TaskSet contains a set of tasks.
type TaskSet map[uint64]*Task

func (ts TaskSet) dup() TaskSet {
	var tasks = TaskSet{}
	for id, task := range ts {
		tasks[id] = task
	}

	return tasks
}

// FilterPriority returns all the tasks with at least the given priority.
func (ts TaskSet) FilterPriority(pri Priority) TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if task.Priority >= pri {
			tasks[id] = task
		}
	}

	return tasks
}

// FilterTag returns all tasks with the given tag.
func (ts TaskSet) FilterTag(tag string) TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if contains(tag, task.Tags) {
			tasks[id] = task
		}
	}

	return tasks
}

// FilterTags returns all tasks with the given tags.
func (ts TaskSet) FilterTags(tags []string) TaskSet {
	var tasks = ts.dup()
	for _, tag := range tags {
		tasks = tasks.FilterTag(tag)
	}

	return tasks
}

// Unfinished returns the subset of tasks that aren't completed.
func (ts TaskSet) Unfinished() TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if !task.Done {
			tasks[id] = task
		}
	}

	return tasks
}

// CompletedDuration returns the tasks completed within the last
// duration, selecting on the completion date.
func (ts TaskSet) CompletedDuration(dur time.Duration) TaskSet {
	var tasks = TaskSet{}

	started := time.Now().Add(-1 * dur)
	for id, task := range ts {
		if task.Done {
			if task.Finished.After(started) {
				tasks[id] = task
			}
		}
	}

	return tasks
}

// CreatedDuration returns the tasks completed within the last
// duration, selecting on the creation date.
func (ts TaskSet) CreatedDuration(dur time.Duration) TaskSet {
	var tasks = TaskSet{}

	started := time.Now().Add(-1 * dur)
	for id, task := range ts {
		if task.Done {
			if task.Created.After(started) {
				tasks[id] = task
			}
		}
	}

	return tasks
}

// CompletedRange returns a list of tasks completed within the
// specified times, selecting on the completion date.
func (ts TaskSet) CompletedRange(start, end time.Time) TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if task.Done {
			if after(task.Finished, start) && before(task.Finished, end) {
				tasks[id] = task
			}
		}
	}

	return tasks
}

// CreatedRange returns a list of tasks completed within the
// specified times, selecting on the created date.
func (ts TaskSet) CreatedRange(start, end time.Time) TaskSet {
	var tasks = TaskSet{}

	for id, task := range ts {
		if task.Done {
			if task.Created.After(start) && task.Created.Before(end) {
				tasks[id] = task
			}
		}
	}

	return tasks
}

// Sort returns a list of the tasks in chronological order.
func (ts TaskSet) Sort() []*Task {
	var ids = make([]uint64, 0, len(ts))
	for id := range ts {
		ids = append(ids, id)
	}

	var tasks = make([]*Task, len(ids))
	for i := range ids {
		tasks[i] = ts[ids[i]]
	}

	return tasks
}

// NewTaskID returns a new task identifier.
func NewTaskID() uint64 {
	return uint64(time.Now().UnixNano())
}
<|endoftext|>


<commit_before>/* RPC with a pool of providers each connected via a web-socket. */
package wsrpcpool

// Main testing module

import (
	"errors"
	"io"
	"net/rpc"
	"testing"
	"sync"
)

/* TestNewPool tests the NewPool works as expected. */
func TestNewPool(t *testing.T) {
	pool := NewPool()
	if pool == nil {
		t.Fail()
	}
}

/* TestNewPoolTLS tests the NewPoolTLS works as expected. */
func TestNewPoolTLS(t *testing.T) {
	pool, err := NewPoolTLS("testfiles/server.crt", "testfiles/server.key")
	if err != nil {
		t.Error(err)
	}
	if pool == nil {
		t.Fatal("Pool is nil")
	}
	checkPoolCerts(t, pool, 1, 0)
}

/* TestNewPoolTLSAuth tests the NewPoolTLSAuth works as expected. */
func TestNewPoolTLSAuth(t *testing.T) {
	pool, err := NewPoolTLSAuth("testfiles/server.crt", "testfiles/server.key", "testfiles/rootCA.crt")
	if err != nil {
		t.Error(err)
	}
	if pool == nil {
		t.Fatal("Pool is nil")
	}
	checkPoolCerts(t, pool, 1, 1)
}

/* checkPoolCerts checks that a specified number of server and root CA
certificates (for client verification) are loaded into the TLSConfig of
the pool server. */
func checkPoolCerts(t *testing.T, pool *PoolServer, certs, clientCAs int) {
	if pool.Server.TLSConfig == nil {
		t.Fatal("TLSConfig is nil")
	}
	csLen := len(pool.Server.TLSConfig.Certificates)
	if csLen != certs {
		t.Errorf("Expected %d certificate, got %d\n", certs, csLen)
	}
	if pool.Server.TLSConfig.ClientCAs == nil {
		if clientCAs > 0 {
			t.Fatal("Client CA pool is nil")
		}
	} else {
		subjs := pool.Server.TLSConfig.ClientCAs.Subjects()
		if len(subjs) != clientCAs {
			t.Errorf("Expected %d certificate in the client CA pool, got %d\n", clientCAs, len(subjs))
		}
	}
}

/* TestBindDefault tests Bind("/") works on a plain pool. */
func TestBindDefault(t *testing.T) {
	pool := NewPool()
	if pool == nil {
		t.Fatal("Unable to instantiate a plain pool")
	}
	pool.Bind("/")
	if pool.DefaultPool == nil {
		t.Fatal("DefaultPool is nil")
	}
}

/* TestBindName tests Bind("/") works on a plain pool for a named provider. */
func TestBindName(t *testing.T) {
	pool := NewPool()
	if pool == nil {
		t.Fatal("Unable to instantiate a plain pool")
	}
	pool.Bind("/", "Test")
	if pool.DefaultPool != nil {
		t.Fatal("DefaultPool is not nil")
	}
	if pool.PoolMap == nil {
		t.Fatal("PoolMap[\"Test\"] is nil")
	}
}

/* closeConns closes the given pool and provider connections. */
func closeConns(t *testing.T, pool io.Closer, closers ...io.Closer) {
	for _, cl := range closers {
		if cl != nil {
			if err := cl.Close(); err != nil {
				t.Error(err)
			}
		}
	}
	if pool != nil {
		if err := pool.Close(); err != nil {
			t.Error(err)
		}
	}
}

/* waitConnected waits for the given connection to either connect
or be closed. Returns the connected flag value. */
func waitConnected(pc *PoolConn) bool {
	var connected bool
	select {
	case <-pc.Connected:
		connected = true
	case <-pc.Closed:
	}
	return connected
}

/* tryConnect tries to connect the given provider making a given number of
attempts at most. Returns the new connection and the connected flag. */
func tryConnect(p *Provider, maxattempts int) (*PoolConn, bool) {
	p.MaxAttempts = maxattempts
	pc := p.ConnectAndServe()
	connected := waitConnected(pc)
	return pc, connected
}

/* tryConnectCaller tries to connect the given provider as a caller
making a given number of attempts at most. Returns the new caller connection
and the connected flag. */
func tryConnectCaller(p *Provider, maxattempts int) (*PoolCallerConn, bool) {
	p.MaxAttempts = maxattempts
	pc := p.ConnectAndUse()
	connected := waitConnected(pc.PoolConn)
	return pc, connected
}

/* testConnection tests for a successful provider to pool server connection
using the given callback functions. */
func testConnection(t *testing.T, getpool func() (*PoolServer, error), listen func(pool *PoolServer) error, getproviders func() ([]*Provider, error), connect func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error)) {
	pool, err := getpool()
	if err != nil {
		t.Fatal(err)
	}
	if pool == nil {
		t.Fatal("Unable to instantiate a plain pool")
	}

	go func() {
		if err := listen(pool); err != nil {
			t.Error(err)
		}
	}()
	<-pool.Listening

	ps, err := getproviders()
	if err != nil {
		t.Error(err)
	}
	if len(ps) == 0 {
		t.Error("Unable to instantiate the provider")
	}

	conns, err := connect(pool, ps...)
	if err != nil {
		t.Error(err)
	}

	closeConns(t, pool, conns...)
}

/* TestConnection tests for a successful provider to pool server connection
over an unencrypted channel. */
func TestConnection(t *testing.T) {
	testConnection(t,
		func() (*PoolServer, error) {
			return NewPool(), nil
		},
		func(pool *PoolServer) error {
			pool.Bind("/")
			return pool.ListenAndUse("localhost:8080")
		},
		func() ([]*Provider, error) {
			p, err := NewProvider("ws://localhost:8080/")
			return []*Provider{p}, err
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			}
			return []io.Closer{pc}, nil
		})
}

/* TestConnectionTLS tests for a successful provider to pool server connection
over an encrypted channel. */
func TestConnectionTLS(t *testing.T) {
	testConnection(t,
		func() (*PoolServer, error) {
			return NewPoolTLS("testfiles/server.crt", "testfiles/server.key")
		},
		func(pool *PoolServer) error {
			pool.Bind("/")
			return pool.ListenAndUseTLS("localhost:8443")
		},
		func() ([]*Provider, error) {
			p, err := NewProvider("wss://localhost:8443/", "testfiles/rootCA.crt")
			return []*Provider{p}, err
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			}
			return []io.Closer{pc}, nil
		})
}

/* TestConnectionTLS expects unsuccessful provider to pool server connection
over an encrypted channel due to unknown server certificate. */
func TestConnectionTLSFail(t *testing.T) {
	testConnection(t,
		func() (*PoolServer, error) {
			return NewPoolTLS("testfiles/server.crt", "testfiles/server.key")
		},
		func(pool *PoolServer) error {
			pool.Bind("/")
			return pool.ListenAndUseTLS("localhost:8443")
		},
		func() ([]*Provider, error) {
			p, err := NewProvider("wss://localhost:8443/")
			return []*Provider{p}, err
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if connected {
				t.Error("Unexpected connection")
			}
			pc.Close()
			return nil, nil
		})
}

/* newTLSAuthProvider return a new provider with root and client
certificates configured. */
func newTLSAuthProvider() (*Provider, error) {
	return NewProviderTLSAuth("wss://localhost:8443/", "testfiles/client.crt", "testfiles/client.key", "testfiles/rootCA.crt")
}

/* testConnectionTLSAuth tests for a successful provider to pool server connection
over an encrypted channel with client-side certificate authentication
using the provided test function and optional hooks. */
func testConnectionTLSAuth(t *testing.T, setupPool func(pool *PoolServer), getproviders func() ([]*Provider, error), connect func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error)) {
	testConnection(t,
		func() (*PoolServer, error) {
			return NewPoolTLSAuth("testfiles/server.crt", "testfiles/server.key", "testfiles/rootCA.crt")
		},
		func(pool *PoolServer) error {
			if setupPool != nil {
				setupPool(pool)
			} else {
				pool.Bind("/")
			}
			return pool.ListenAndUseTLS("localhost:8443")
		},
		func() ([]*Provider, error) {
			if getproviders != nil {
				return getproviders()
			} else {
				p, err := newTLSAuthProvider()
				return []*Provider{p}, err
			}
		},
		connect)
}

/* TestConnectionTLSAuth tests for a successful provider to pool server connection
over an encrypted channel with client-side certificate authentication. */
func TestConnectionTLSAuth(t *testing.T) {
	testConnectionTLSAuth(t, nil, nil,
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			}
			return []io.Closer{pc}, nil
		})
}

/* TestConnectionTLSAuthFail tests for a *unsuccessful* provider to pool server
connection over an encrypted channel without the expected client-side certificate
authentication. */
func TestConnectionTLSAuthFail(t *testing.T) {
	testConnectionTLSAuth(t, nil,
		func() ([]*Provider, error) {
			p, err := NewProvider("wss://localhost:8443/", "testfiles/rootCA.crt")
			return []*Provider{p}, err
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if connected {
				t.Error("Unexpected connection")
			}
			pc.Close()
			return nil, nil
		})
}

/* Test is used as the exported function provider for remote calls. */
type Test struct {
}

/* TestFunc sets its return argument to the value of its first argument. */
func (t *Test) TestFunc(arg int, reply *int) error {
	*reply = arg
	return nil
}

/* TestError returns the error value with the given text. It also sets
its return value to the same string if the pointer is not nil. */
func (t *Test) TestError(text string, reply *string) error {
	if reply != nil {
		*reply = text
	}
	return errors.New(text)
}

/* rpcClient provides Go and Call methods. */
type rpcClient interface {
	Go(serviceMethod string, args interface{}, reply interface{}, done chan *rpc.Call) (*rpc.Call, error)
	Call(serviceMethod string, args interface{}, reply interface{}) error
}

/* testCalls perform the series of remote calls using the Test type. */
func testCalls(t *testing.T, client rpcClient) {
	rpc.Register(&Test{})

	var testval int = 1
	var reply int
	if err := client.Call("Test.TestFunc", testval, &reply); err != nil {
		t.Error(err)
	}
	if reply != testval {
		t.Error("Unexpected result")
	}

	var errtext = "Error text"
	if err := client.Call("Test.TestError", errtext, nil); err == nil {
		t.Error("Expected error")
	} else {
		if err.Error() != errtext {
			t.Error(err)
		}
	}
}

/* TestCallTLSAuth tests for a successful method call over
an encrypted channel with client-side certificate authentication. */
func TestCallTLSAuth(t *testing.T) {
	testConnectionTLSAuth(t, nil, nil,
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			} else {
				testCalls(t, pool)
			}
			return []io.Closer{pc}, nil
		})
}

/* TestInCallTLSAuth tests for a successful incoming method call over
an encrypted channel with client-side certificate authentication. */
func TestInCallTLSAuth(t *testing.T) {
	testConnectionTLSAuth(t,
		func(pool *PoolServer) {
			pool.BindIn("/")
		},
		nil,
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnectCaller(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			} else {
				testCalls(t, pc)
			}
			return []io.Closer{pc}, nil
		})
}

var (
	/* A number of providers */
	ProviderCount int = 5
	/* A number of connection per provider */
	ConnsPerProvider int = 5
)

/* TestConnectionTLSAuthMulti tests for a successful connection between the pool
server and the set of providers over an encrypted channel with client-side certificate
authentication. */
func TestConnectionTLSAuthMulti(t *testing.T) {
	testConnectionTLSAuth(t, nil,
		func() ([]*Provider, error) {
			ps := make([]*Provider, 0, ProviderCount)
			for i := 0; i < ProviderCount; i++ {
				if p, err := newTLSAuthProvider(); err == nil {
					ps = append(ps, p)
				} else {
					return nil, err
				}
			}
			return ps, nil
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			wg := sync.WaitGroup{}
			pcs := make(chan io.Closer, len(ps) * ConnsPerProvider)
			for _, p := range ps {
				for i := 0; i < ConnsPerProvider; i++ {
					go func() {
						wg.Add(1)
						pc, connected := tryConnect(p, 1)
						if !connected {
							t.Error("Not connected")
						}
						pcs <- pc
						wg.Done()
					}()
				}
			}
			wg.Wait()
			close(pcs)

			return func() []io.Closer {
				cls := make([]io.Closer, 0, len(ps) * ConnsPerProvider)
				for pc := range(pcs) {
					cls = append(cls, pc)
				}
				return cls
			}(), nil
		})
}
<commit_msg>Fix: Add items to the wait group *synchronously*<commit_after>/* RPC with a pool of providers each connected via a web-socket. */
package wsrpcpool

// Main testing module

import (
	"errors"
	"io"
	"net/rpc"
	"testing"
	"sync"
)

/* TestNewPool tests the NewPool works as expected. */
func TestNewPool(t *testing.T) {
	pool := NewPool()
	if pool == nil {
		t.Fail()
	}
}

/* TestNewPoolTLS tests the NewPoolTLS works as expected. */
func TestNewPoolTLS(t *testing.T) {
	pool, err := NewPoolTLS("testfiles/server.crt", "testfiles/server.key")
	if err != nil {
		t.Error(err)
	}
	if pool == nil {
		t.Fatal("Pool is nil")
	}
	checkPoolCerts(t, pool, 1, 0)
}

/* TestNewPoolTLSAuth tests the NewPoolTLSAuth works as expected. */
func TestNewPoolTLSAuth(t *testing.T) {
	pool, err := NewPoolTLSAuth("testfiles/server.crt", "testfiles/server.key", "testfiles/rootCA.crt")
	if err != nil {
		t.Error(err)
	}
	if pool == nil {
		t.Fatal("Pool is nil")
	}
	checkPoolCerts(t, pool, 1, 1)
}

/* checkPoolCerts checks that a specified number of server and root CA
certificates (for client verification) are loaded into the TLSConfig of
the pool server. */
func checkPoolCerts(t *testing.T, pool *PoolServer, certs, clientCAs int) {
	if pool.Server.TLSConfig == nil {
		t.Fatal("TLSConfig is nil")
	}
	csLen := len(pool.Server.TLSConfig.Certificates)
	if csLen != certs {
		t.Errorf("Expected %d certificate, got %d\n", certs, csLen)
	}
	if pool.Server.TLSConfig.ClientCAs == nil {
		if clientCAs > 0 {
			t.Fatal("Client CA pool is nil")
		}
	} else {
		subjs := pool.Server.TLSConfig.ClientCAs.Subjects()
		if len(subjs) != clientCAs {
			t.Errorf("Expected %d certificate in the client CA pool, got %d\n", clientCAs, len(subjs))
		}
	}
}

/* TestBindDefault tests Bind("/") works on a plain pool. */
func TestBindDefault(t *testing.T) {
	pool := NewPool()
	if pool == nil {
		t.Fatal("Unable to instantiate a plain pool")
	}
	pool.Bind("/")
	if pool.DefaultPool == nil {
		t.Fatal("DefaultPool is nil")
	}
}

/* TestBindName tests Bind("/") works on a plain pool for a named provider. */
func TestBindName(t *testing.T) {
	pool := NewPool()
	if pool == nil {
		t.Fatal("Unable to instantiate a plain pool")
	}
	pool.Bind("/", "Test")
	if pool.DefaultPool != nil {
		t.Fatal("DefaultPool is not nil")
	}
	if pool.PoolMap == nil {
		t.Fatal("PoolMap[\"Test\"] is nil")
	}
}

/* closeConns closes the given pool and provider connections. */
func closeConns(t *testing.T, pool io.Closer, closers ...io.Closer) {
	for _, cl := range closers {
		if cl != nil {
			if err := cl.Close(); err != nil {
				t.Error(err)
			}
		}
	}
	if pool != nil {
		if err := pool.Close(); err != nil {
			t.Error(err)
		}
	}
}

/* waitConnected waits for the given connection to either connect
or be closed. Returns the connected flag value. */
func waitConnected(pc *PoolConn) bool {
	var connected bool
	select {
	case <-pc.Connected:
		connected = true
	case <-pc.Closed:
	}
	return connected
}

/* tryConnect tries to connect the given provider making a given number of
attempts at most. Returns the new connection and the connected flag. */
func tryConnect(p *Provider, maxattempts int) (*PoolConn, bool) {
	p.MaxAttempts = maxattempts
	pc := p.ConnectAndServe()
	connected := waitConnected(pc)
	return pc, connected
}

/* tryConnectCaller tries to connect the given provider as a caller
making a given number of attempts at most. Returns the new caller connection
and the connected flag. */
func tryConnectCaller(p *Provider, maxattempts int) (*PoolCallerConn, bool) {
	p.MaxAttempts = maxattempts
	pc := p.ConnectAndUse()
	connected := waitConnected(pc.PoolConn)
	return pc, connected
}

/* testConnection tests for a successful provider to pool server connection
using the given callback functions. */
func testConnection(t *testing.T, getpool func() (*PoolServer, error), listen func(pool *PoolServer) error, getproviders func() ([]*Provider, error), connect func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error)) {
	pool, err := getpool()
	if err != nil {
		t.Fatal(err)
	}
	if pool == nil {
		t.Fatal("Unable to instantiate a plain pool")
	}

	go func() {
		if err := listen(pool); err != nil {
			t.Error(err)
		}
	}()
	<-pool.Listening

	ps, err := getproviders()
	if err != nil {
		t.Error(err)
	}
	if len(ps) == 0 {
		t.Error("Unable to instantiate the provider")
	}

	conns, err := connect(pool, ps...)
	if err != nil {
		t.Error(err)
	}

	closeConns(t, pool, conns...)
}

/* TestConnection tests for a successful provider to pool server connection
over an unencrypted channel. */
func TestConnection(t *testing.T) {
	testConnection(t,
		func() (*PoolServer, error) {
			return NewPool(), nil
		},
		func(pool *PoolServer) error {
			pool.Bind("/")
			return pool.ListenAndUse("localhost:8080")
		},
		func() ([]*Provider, error) {
			p, err := NewProvider("ws://localhost:8080/")
			return []*Provider{p}, err
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			}
			return []io.Closer{pc}, nil
		})
}

/* TestConnectionTLS tests for a successful provider to pool server connection
over an encrypted channel. */
func TestConnectionTLS(t *testing.T) {
	testConnection(t,
		func() (*PoolServer, error) {
			return NewPoolTLS("testfiles/server.crt", "testfiles/server.key")
		},
		func(pool *PoolServer) error {
			pool.Bind("/")
			return pool.ListenAndUseTLS("localhost:8443")
		},
		func() ([]*Provider, error) {
			p, err := NewProvider("wss://localhost:8443/", "testfiles/rootCA.crt")
			return []*Provider{p}, err
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			}
			return []io.Closer{pc}, nil
		})
}

/* TestConnectionTLS expects unsuccessful provider to pool server connection
over an encrypted channel due to unknown server certificate. */
func TestConnectionTLSFail(t *testing.T) {
	testConnection(t,
		func() (*PoolServer, error) {
			return NewPoolTLS("testfiles/server.crt", "testfiles/server.key")
		},
		func(pool *PoolServer) error {
			pool.Bind("/")
			return pool.ListenAndUseTLS("localhost:8443")
		},
		func() ([]*Provider, error) {
			p, err := NewProvider("wss://localhost:8443/")
			return []*Provider{p}, err
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if connected {
				t.Error("Unexpected connection")
			}
			pc.Close()
			return nil, nil
		})
}

/* newTLSAuthProvider return a new provider with root and client
certificates configured. */
func newTLSAuthProvider() (*Provider, error) {
	return NewProviderTLSAuth("wss://localhost:8443/", "testfiles/client.crt", "testfiles/client.key", "testfiles/rootCA.crt")
}

/* testConnectionTLSAuth tests for a successful provider to pool server connection
over an encrypted channel with client-side certificate authentication
using the provided test function and optional hooks. */
func testConnectionTLSAuth(t *testing.T, setupPool func(pool *PoolServer), getproviders func() ([]*Provider, error), connect func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error)) {
	testConnection(t,
		func() (*PoolServer, error) {
			return NewPoolTLSAuth("testfiles/server.crt", "testfiles/server.key", "testfiles/rootCA.crt")
		},
		func(pool *PoolServer) error {
			if setupPool != nil {
				setupPool(pool)
			} else {
				pool.Bind("/")
			}
			return pool.ListenAndUseTLS("localhost:8443")
		},
		func() ([]*Provider, error) {
			if getproviders != nil {
				return getproviders()
			} else {
				p, err := newTLSAuthProvider()
				return []*Provider{p}, err
			}
		},
		connect)
}

/* TestConnectionTLSAuth tests for a successful provider to pool server connection
over an encrypted channel with client-side certificate authentication. */
func TestConnectionTLSAuth(t *testing.T) {
	testConnectionTLSAuth(t, nil, nil,
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			}
			return []io.Closer{pc}, nil
		})
}

/* TestConnectionTLSAuthFail tests for a *unsuccessful* provider to pool server
connection over an encrypted channel without the expected client-side certificate
authentication. */
func TestConnectionTLSAuthFail(t *testing.T) {
	testConnectionTLSAuth(t, nil,
		func() ([]*Provider, error) {
			p, err := NewProvider("wss://localhost:8443/", "testfiles/rootCA.crt")
			return []*Provider{p}, err
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if connected {
				t.Error("Unexpected connection")
			}
			pc.Close()
			return nil, nil
		})
}

/* Test is used as the exported function provider for remote calls. */
type Test struct {
}

/* TestFunc sets its return argument to the value of its first argument. */
func (t *Test) TestFunc(arg int, reply *int) error {
	*reply = arg
	return nil
}

/* TestError returns the error value with the given text. It also sets
its return value to the same string if the pointer is not nil. */
func (t *Test) TestError(text string, reply *string) error {
	if reply != nil {
		*reply = text
	}
	return errors.New(text)
}

/* rpcClient provides Go and Call methods. */
type rpcClient interface {
	Go(serviceMethod string, args interface{}, reply interface{}, done chan *rpc.Call) (*rpc.Call, error)
	Call(serviceMethod string, args interface{}, reply interface{}) error
}

/* testCalls perform the series of remote calls using the Test type. */
func testCalls(t *testing.T, client rpcClient) {
	rpc.Register(&Test{})

	var testval int = 1
	var reply int
	if err := client.Call("Test.TestFunc", testval, &reply); err != nil {
		t.Error(err)
	}
	if reply != testval {
		t.Error("Unexpected result")
	}

	var errtext = "Error text"
	if err := client.Call("Test.TestError", errtext, nil); err == nil {
		t.Error("Expected error")
	} else {
		if err.Error() != errtext {
			t.Error(err)
		}
	}
}

/* TestCallTLSAuth tests for a successful method call over
an encrypted channel with client-side certificate authentication. */
func TestCallTLSAuth(t *testing.T) {
	testConnectionTLSAuth(t, nil, nil,
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnect(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			} else {
				testCalls(t, pool)
			}
			return []io.Closer{pc}, nil
		})
}

/* TestInCallTLSAuth tests for a successful incoming method call over
an encrypted channel with client-side certificate authentication. */
func TestInCallTLSAuth(t *testing.T) {
	testConnectionTLSAuth(t,
		func(pool *PoolServer) {
			pool.BindIn("/")
		},
		nil,
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			pc, connected := tryConnectCaller(ps[0], 1)
			if !connected {
				t.Error("Not connected")
			} else {
				testCalls(t, pc)
			}
			return []io.Closer{pc}, nil
		})
}

var (
	/* A number of providers */
	ProviderCount int = 5
	/* A number of connection per provider */
	ConnsPerProvider int = 5
)

/* TestConnectionTLSAuthMulti tests for a successful connection between the pool
server and the set of providers over an encrypted channel with client-side certificate
authentication. */
func TestConnectionTLSAuthMulti(t *testing.T) {
	testConnectionTLSAuth(t, nil,
		func() ([]*Provider, error) {
			ps := make([]*Provider, 0, ProviderCount)
			for i := 0; i < ProviderCount; i++ {
				if p, err := newTLSAuthProvider(); err == nil {
					ps = append(ps, p)
				} else {
					return nil, err
				}
			}
			return ps, nil
		},
		func(pool *PoolServer, ps ...*Provider) ([]io.Closer, error) {
			wg := sync.WaitGroup{}
			pcs := make(chan io.Closer, len(ps) * ConnsPerProvider)
			for _, p := range ps {
				for i := 0; i < ConnsPerProvider; i++ {
					wg.Add(1)
					go func() {
						defer wg.Done()
						pc, connected := tryConnect(p, 1)
						if !connected {
							t.Error("Not connected")
						}
						pcs <- pc
					}()
				}
			}
			wg.Wait()
			close(pcs)

			return func() []io.Closer {
				cls := make([]io.Closer, 0, len(ps) * ConnsPerProvider)
				for pc := range pcs {
					cls = append(cls, pc)
				}
				return cls
			}(), nil
		})
}
<|endoftext|>


<commit_before>package config

import (
	"errors"
	"os"
	"strings"

	"github.com/spf13/viper"
)

func _getAuth(section string) (string, string, error) {
	if os.Getenv("ONDEVICE_USER") != "" || os.Getenv("ONDEVICE_AUTH") != "" {
		return os.Getenv("ONDEVICE_USER"), os.Getenv("ONDEVICE_AUTH"), nil
	}

	username, uerr := GetString(section, "user")
	auth, aerr := GetString(section, "auth")

	if uerr != nil {
		return "", "", uerr
	}
	if aerr != nil {
		return "", "", aerr
	}

	return username, auth, nil
}

// GetClientAuth -- Get the default client authentication information
func GetClientAuth() (string, string, error) {
	return _getAuth("client")
}

// GetClientUserAuth -- get the authentication credentials for a specific client user
func GetClientUserAuth(username string) (string, string, error) {
	defaultU, defaultA, defaultE := GetClientAuth()
	if defaultE == nil && defaultU == username {
		return defaultU, defaultA, nil
	}

	auth, err := GetString("client", "auth_"+username)
	if err == nil {
		return username, auth, nil
	}

	return "", "", errors.New("no client credentials found for user")
}

// GetDeviceAuth -- Get the device authentication
func GetDeviceAuth() (string, string, error) {
	return _getAuth("device")
}

// ListAuthenticatedUsers -- returns the names of users we have client auth for
func ListAuthenticatedUsers() []string {
	// TODO this is messy but will do for now - we'll improve this once we have a separate auth file

	var rc []string
	var uniqueUsers = make(map[string]bool)

	if mainUser := viper.GetString("client.user"); mainUser != "" {
		rc = append(rc, mainUser)
		uniqueUsers[strings.ToLower(mainUser)] = true
	}

	for _, k := range viper.AllKeys() {
		if strings.HasPrefix(k, "client.auth_") {
			var name = k[12:]
			var lowerName = strings.ToLower(name)
			if !uniqueUsers[lowerName] {
				rc = append(rc, k[12:])
				uniqueUsers[lowerName] = true
			}
		}
	}

	return rc
}
<commit_msg>config.ListAuthenticatedUsers(): replaced viper with AllValues()<commit_after>package config

import (
	"errors"
	"os"
	"strings"

	"github.com/sirupsen/logrus"
)

func _getAuth(section string) (string, string, error) {
	if os.Getenv("ONDEVICE_USER") != "" || os.Getenv("ONDEVICE_AUTH") != "" {
		return os.Getenv("ONDEVICE_USER"), os.Getenv("ONDEVICE_AUTH"), nil
	}

	username, uerr := GetString(section, "user")
	auth, aerr := GetString(section, "auth")

	if uerr != nil {
		return "", "", uerr
	}
	if aerr != nil {
		return "", "", aerr
	}

	return username, auth, nil
}

// GetClientAuth -- Get the default client authentication information
func GetClientAuth() (string, string, error) {
	return _getAuth("client")
}

// GetClientUserAuth -- get the authentication credentials for a specific client user
func GetClientUserAuth(username string) (string, string, error) {
	defaultU, defaultA, defaultE := GetClientAuth()
	if defaultE == nil && defaultU == username {
		return defaultU, defaultA, nil
	}

	auth, err := GetString("client", "auth_"+username)
	if err == nil {
		return username, auth, nil
	}

	return "", "", errors.New("no client credentials found for user")
}

// GetDeviceAuth -- Get the device authentication
func GetDeviceAuth() (string, string, error) {
	return _getAuth("device")
}

// ListAuthenticatedUsers -- returns the names of users we have client auth for
func ListAuthenticatedUsers() []string {
	// TODO this is messy but will do for now - we'll improve this once we have a separate auth file

	var rc []string
	var uniqueUsers = make(map[string]bool)

	if mainUser, err := GetString("client", "user"); err != nil && mainUser != "" {
		rc = append(rc, mainUser)
		uniqueUsers[strings.ToLower(mainUser)] = true
	}

	var cfg, err = AllValues()
	if err != nil {
		logrus.WithError(err).Fatal("failed to fetch configuration")
	}
	for k := range cfg {
		if strings.HasPrefix(k, "client.auth_") {
			var name = k[12:]
			var lowerName = strings.ToLower(name)
			if !uniqueUsers[lowerName] {
				rc = append(rc, k[12:])
				uniqueUsers[lowerName] = true
			}
		}
	}

	return rc
}
<|endoftext|>