<commit_before>package ecs

import (
	"time"

	"github.com/denverdino/aliyungo/util"
)

// InstanceStatus represents instance status
type InstanceStatus string

// Constants of InstanceStatus
const (
	Creating = InstanceStatus("Creating")
	Running  = InstanceStatus("Running")
	Starting = InstanceStatus("Starting")

	Stopped  = InstanceStatus("Stopped")
	Stopping = InstanceStatus("Stopping")
)

type InternetChargeType string

const (
	PayByBandwidth = InternetChargeType("PayByBandwidth")
	PayByTraffic   = InternetChargeType("PayByTraffic")
)

type LockReason string

const (
	LockReasonFinancial = LockReason("financial")
	LockReasonSecurity  = LockReason("security")
)

type DescribeInstanceStatusArgs struct {
	RegionId Region
	ZoneId   string
	Pagination
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instancestatusitemtype
type InstanceStatusItemType struct {
	InstanceId string
	Status     InstanceStatus
}

type DescribeInstanceStatusResponse struct {
	CommonResponse
	PaginationResult
	InstanceStatuses struct {
		InstanceStatus []InstanceStatusItemType
	}
}

// DescribeInstanceStatus describes instance status
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancestatus
func (client *Client) DescribeInstanceStatus(args *DescribeInstanceStatusArgs) (instanceStatuses []InstanceStatusItemType, pagination *PaginationResult, err error) {
	args.validate()
	response := DescribeInstanceStatusResponse{}

	err = client.Invoke("DescribeInstanceStatus", args, &response)

	if err == nil {
		return response.InstanceStatuses.InstanceStatus, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type StopInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type StopInstanceResponse struct {
	CommonResponse
}

// StopInstance stops instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&stopinstance
func (client *Client) StopInstance(instanceId string, forceStop bool) error {
	args := StopInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := StopInstanceResponse{}
	err := client.Invoke("StopInstance", &args, &response)
	return err
}

type StartInstanceArgs struct {
	InstanceId string
}

type StartInstanceResponse struct {
	CommonResponse
}

// StartInstance starts instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&startinstance
func (client *Client) StartInstance(instanceId string) error {
	args := StartInstanceArgs{InstanceId: instanceId}
	response := StartInstanceResponse{}
	err := client.Invoke("StartInstance", &args, &response)
	return err
}

type RebootInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type RebootInstanceResponse struct {
	CommonResponse
}

// RebootInstance reboot instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&rebootinstance
func (client *Client) RebootInstance(instanceId string, forceStop bool) error {
	request := RebootInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := RebootInstanceResponse{}
	err := client.Invoke("RebootInstance", &request, &response)
	return err
}

type DescribeInstanceAttributeArgs struct {
	InstanceId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&operationlockstype
type OperationLocksType struct {
	LockReason []LockReason //enum for financial, security
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&securitygroupidsettype
type SecurityGroupIdSetType struct {
	SecurityGroupId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&ipaddresssettype
type IpAddressSetType struct {
	IpAddress []string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&vpcattributestype
type VpcAttributesType struct {
	VpcId            string
	VSwitchId        string
	PrivateIpAddress IpAddressSetType
	NatIpAddress     string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&eipaddressassociatetype
type EipAddressAssociateType struct {
	AllocationId       string
	IpAddress          string
	Bandwidth          int
	InternetChargeType InternetChargeType
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instanceattributestype
type InstanceAttributesType struct {
	InstanceId       string
	InstanceName     string
	Description      string
	ImageId          string
	RegionId         Region
	ZoneId           string
	ClusterId        string
	InstanceType     string
	HostName         string
	Status           InstanceStatus
	OperationLocks   OperationLocksType
	SecurityGroupIds struct {
		SecurityGroupId []string
	}
	PublicIpAddress         IpAddressSetType
	InnerIpAddress          IpAddressSetType
	InstanceNetworkType     string //enum Classic | Vpc
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	InternetChargeType      InternetChargeType
	CreationTime            util.ISO6801Time //time.Time
	VpcAttributes           VpcAttributesType
	EipAddress              EipAddressAssociateType
}

type DescribeInstanceAttributeResponse struct {
	CommonResponse
	InstanceAttributesType
}

// DescribeInstanceAttribute describes instance attribute
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstanceattribute
func (client *Client) DescribeInstanceAttribute(instanceId string) (instance *InstanceAttributesType, err error) {
	args := DescribeInstanceAttributeArgs{InstanceId: instanceId}

	response := DescribeInstanceAttributeResponse{}
	err = client.Invoke("DescribeInstanceAttribute", &args, &response)
	if err != nil {
		return nil, err
	}
	return &response.InstanceAttributesType, err
}

// Default timeout value for WaitForInstance method
const InstanceDefaultTimeout = 120

// WaitForInstance waits for instance to given status
func (client *Client) WaitForInstance(instanceId string, status InstanceStatus, timeout int) error {
	if timeout <= 0 {
		timeout = InstanceDefaultTimeout
	}
	for {
		instance, err := client.DescribeInstanceAttribute(instanceId)
		if err != nil {
			return err
		}
		if instance.Status == status {
			//TODO
			//Sleep one more time for timing issues
			time.Sleep(DefaultWaitForInterval * time.Second)
			break
		}
		timeout = timeout - DefaultWaitForInterval
		if timeout <= 0 {
			return getECSErrorFromString("Timeout")
		}
		time.Sleep(DefaultWaitForInterval * time.Second)

	}
	return nil
}

type DescribeInstanceVncUrlArgs struct {
	RegionId   Region
	InstanceId string
}

type DescribeInstanceVncUrlResponse struct {
	CommonResponse
	VncUrl string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancevncurl
func (client *Client) DescribeInstanceVncUrl(args *DescribeInstanceVncUrlArgs) (string, error) {
	response := DescribeInstanceVncUrlResponse{}

	err := client.Invoke("DescribeInstanceVncUrl", args, &response)

	if err == nil {
		return response.VncUrl, nil
	}

	return "", err
}

type DescribeInstancesArgs struct {
	RegionId            Region
	VpcId               string
	VSwitchId           string
	ZoneId              string
	InstanceIds         string
	InstanceNetworkType string
	PrivateIpAddresses  string
	InnerIpAddresses    string
	PublicIpAddresses   string
	SecurityGroupId     string
	Pagination
}

type DescribeInstancesResponse struct {
	CommonResponse
	PaginationResult
	Instances struct {
		Instance []InstanceAttributesType
	}
}

// DescribeInstances describes instances
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstances
func (client *Client) DescribeInstances(args *DescribeInstancesArgs) (instances []InstanceAttributesType, pagination *PaginationResult, err error) {
	args.validate()
	response := DescribeInstancesResponse{}

	err = client.Invoke("DescribeInstances", args, &response)

	if err == nil {
		return response.Instances.Instance, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type DeleteInstanceArgs struct {
	InstanceId string
}

type DeleteInstanceResponse struct {
	CommonResponse
}

// DeleteInstance deletes instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&deleteinstance
func (client *Client) DeleteInstance(instanceId string) error {
	args := DeleteInstanceArgs{InstanceId: instanceId}
	response := DeleteInstanceResponse{}
	err := client.Invoke("DeleteInstance", &args, &response)
	return err
}

type DataDiskType struct {
	Size               int
	Category           DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	SnapshotId         string
	DiskName           string
	Description        string
	Device             string
	DeleteWithInstance bool
}

type SystemDiskType struct {
	Category    DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	DiskName    string
	Description string
}

type CreateInstanceArgs struct {
	RegionId                Region
	ZoneId                  string
	ImageId                 string
	InstanceType            string
	SecurityGroupId         string
	InstanceName            string
	Description             string
	InternetChargeType      InternetChargeType
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	HostName                string
	Password                string
	SystemDisk              SystemDiskType
	DataDisk                []DataDiskType
	VSwitchId               string
	PrivateIpAddress        string
	ClientToken             string
}

type CreateInstanceResponse struct {
	CommonResponse
	InstanceId string
}

// CreateInstance creates instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&createinstance
func (client *Client) CreateInstance(args *CreateInstanceArgs) (instanceId string, err error) {
	response := CreateInstanceResponse{}
	err = client.Invoke("CreateInstance", args, &response)
	if err != nil {
		return "", err
	}
	return response.InstanceId, err
}
<commit_msg>debug<commit_after>package ecs

import (
	"time"
)

// InstanceStatus represents instance status
type InstanceStatus string

// Constants of InstanceStatus
const (
	Creating = InstanceStatus("Creating")
	Running  = InstanceStatus("Running")
	Starting = InstanceStatus("Starting")

	Stopped  = InstanceStatus("Stopped")
	Stopping = InstanceStatus("Stopping")
)

type InternetChargeType string

const (
	PayByBandwidth = InternetChargeType("PayByBandwidth")
	PayByTraffic   = InternetChargeType("PayByTraffic")
)

type LockReason string

const (
	LockReasonFinancial = LockReason("financial")
	LockReasonSecurity  = LockReason("security")
)

type DescribeInstanceStatusArgs struct {
	RegionId Region
	ZoneId   string
	Pagination
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instancestatusitemtype
type InstanceStatusItemType struct {
	InstanceId string
	Status     InstanceStatus
}

type DescribeInstanceStatusResponse struct {
	CommonResponse
	PaginationResult
	InstanceStatuses struct {
		InstanceStatus []InstanceStatusItemType
	}
}

// DescribeInstanceStatus describes instance status
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancestatus
func (client *Client) DescribeInstanceStatus(args *DescribeInstanceStatusArgs) (instanceStatuses []InstanceStatusItemType, pagination *PaginationResult, err error) {
	args.validate()
	response := DescribeInstanceStatusResponse{}

	err = client.Invoke("DescribeInstanceStatus", args, &response)

	if err == nil {
		return response.InstanceStatuses.InstanceStatus, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type StopInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type StopInstanceResponse struct {
	CommonResponse
}

// StopInstance stops instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&stopinstance
func (client *Client) StopInstance(instanceId string, forceStop bool) error {
	args := StopInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := StopInstanceResponse{}
	err := client.Invoke("StopInstance", &args, &response)
	return err
}

type StartInstanceArgs struct {
	InstanceId string
}

type StartInstanceResponse struct {
	CommonResponse
}

// StartInstance starts instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&startinstance
func (client *Client) StartInstance(instanceId string) error {
	args := StartInstanceArgs{InstanceId: instanceId}
	response := StartInstanceResponse{}
	err := client.Invoke("StartInstance", &args, &response)
	return err
}

type RebootInstanceArgs struct {
	InstanceId string
	ForceStop  bool
}

type RebootInstanceResponse struct {
	CommonResponse
}

// RebootInstance reboot instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&rebootinstance
func (client *Client) RebootInstance(instanceId string, forceStop bool) error {
	request := RebootInstanceArgs{
		InstanceId: instanceId,
		ForceStop:  forceStop,
	}
	response := RebootInstanceResponse{}
	err := client.Invoke("RebootInstance", &request, &response)
	return err
}

type DescribeInstanceAttributeArgs struct {
	InstanceId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&operationlockstype
type OperationLocksType struct {
	LockReason []LockReason //enum for financial, security
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&securitygroupidsettype
type SecurityGroupIdSetType struct {
	SecurityGroupId string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&ipaddresssettype
type IpAddressSetType struct {
	IpAddress []string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&vpcattributestype
type VpcAttributesType struct {
	VpcId            string
	VSwitchId        string
	PrivateIpAddress IpAddressSetType
	NatIpAddress     string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&eipaddressassociatetype
type EipAddressAssociateType struct {
	AllocationId       string
	IpAddress          string
	Bandwidth          int
	InternetChargeType InternetChargeType
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/datatype&instanceattributestype
type InstanceAttributesType struct {
	InstanceId       string
	InstanceName     string
	Description      string
	ImageId          string
	RegionId         Region
	ZoneId           string
	ClusterId        string
	InstanceType     string
	HostName         string
	Status           InstanceStatus
	OperationLocks   OperationLocksType
	SecurityGroupIds struct {
		SecurityGroupId []string
	}
	PublicIpAddress         IpAddressSetType
	InnerIpAddress          IpAddressSetType
	InstanceNetworkType     string //enum Classic | Vpc
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	InternetChargeType      InternetChargeType
	CreationTime            string
	VpcAttributes           VpcAttributesType
	EipAddress              EipAddressAssociateType
}

type DescribeInstanceAttributeResponse struct {
	CommonResponse
	InstanceAttributesType
}

// DescribeInstanceAttribute describes instance attribute
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstanceattribute
func (client *Client) DescribeInstanceAttribute(instanceId string) (instance *InstanceAttributesType, err error) {
	args := DescribeInstanceAttributeArgs{InstanceId: instanceId}

	response := DescribeInstanceAttributeResponse{}
	err = client.Invoke("DescribeInstanceAttribute", &args, &response)
	if err != nil {
		return nil, err
	}
	return &response.InstanceAttributesType, err
}

// Default timeout value for WaitForInstance method
const InstanceDefaultTimeout = 120

// WaitForInstance waits for instance to given status
func (client *Client) WaitForInstance(instanceId string, status InstanceStatus, timeout int) error {
	if timeout <= 0 {
		timeout = InstanceDefaultTimeout
	}
	for {
		instance, err := client.DescribeInstanceAttribute(instanceId)
		if err != nil {
			return err
		}
		if instance.Status == status {
			//TODO
			//Sleep one more time for timing issues
			time.Sleep(DefaultWaitForInterval * time.Second)
			break
		}
		timeout = timeout - DefaultWaitForInterval
		if timeout <= 0 {
			return getECSErrorFromString("Timeout")
		}
		time.Sleep(DefaultWaitForInterval * time.Second)

	}
	return nil
}

type DescribeInstanceVncUrlArgs struct {
	RegionId   Region
	InstanceId string
}

type DescribeInstanceVncUrlResponse struct {
	CommonResponse
	VncUrl string
}

//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstancevncurl
func (client *Client) DescribeInstanceVncUrl(args *DescribeInstanceVncUrlArgs) (string, error) {
	response := DescribeInstanceVncUrlResponse{}

	err := client.Invoke("DescribeInstanceVncUrl", args, &response)

	if err == nil {
		return response.VncUrl, nil
	}

	return "", err
}

type DescribeInstancesArgs struct {
	RegionId            Region
	VpcId               string
	VSwitchId           string
	ZoneId              string
	InstanceIds         string
	InstanceNetworkType string
	PrivateIpAddresses  string
	InnerIpAddresses    string
	PublicIpAddresses   string
	SecurityGroupId     string
	Pagination
}

type DescribeInstancesResponse struct {
	CommonResponse
	PaginationResult
	Instances struct {
		Instance []InstanceAttributesType
	}
}

// DescribeInstances describes instances
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&describeinstances
func (client *Client) DescribeInstances(args *DescribeInstancesArgs) (instances []InstanceAttributesType, pagination *PaginationResult, err error) {
	args.validate()
	response := DescribeInstancesResponse{}

	err = client.Invoke("DescribeInstances", args, &response)

	if err == nil {
		return response.Instances.Instance, &response.PaginationResult, nil
	}

	return nil, nil, err
}

type DeleteInstanceArgs struct {
	InstanceId string
}

type DeleteInstanceResponse struct {
	CommonResponse
}

// DeleteInstance deletes instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&deleteinstance
func (client *Client) DeleteInstance(instanceId string) error {
	args := DeleteInstanceArgs{InstanceId: instanceId}
	response := DeleteInstanceResponse{}
	err := client.Invoke("DeleteInstance", &args, &response)
	return err
}

type DataDiskType struct {
	Size               int
	Category           DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	SnapshotId         string
	DiskName           string
	Description        string
	Device             string
	DeleteWithInstance bool
}

type SystemDiskType struct {
	Category    DiskCategory //Enum cloud, ephemeral, ephemeral_ssd
	DiskName    string
	Description string
}

type CreateInstanceArgs struct {
	RegionId                Region
	ZoneId                  string
	ImageId                 string
	InstanceType            string
	SecurityGroupId         string
	InstanceName            string
	Description             string
	InternetChargeType      InternetChargeType
	InternetMaxBandwidthIn  int
	InternetMaxBandwidthOut int
	HostName                string
	Password                string
	SystemDisk              SystemDiskType
	DataDisk                []DataDiskType
	VSwitchId               string
	PrivateIpAddress        string
	ClientToken             string
}

type CreateInstanceResponse struct {
	CommonResponse
	InstanceId string
}

// CreateInstance creates instance
//
// You can read doc at http://docs.aliyun.com/#/pub/ecs/open-api/instance&createinstance
func (client *Client) CreateInstance(args *CreateInstanceArgs) (instanceId string, err error) {
	response := CreateInstanceResponse{}
	err = client.Invoke("CreateInstance", args, &response)
	if err != nil {
		return "", err
	}
	return response.InstanceId, err
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package endpoints

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"reflect"
	"strings"

	"golang.org/x/net/context"
	// Mainly for debug logging
	"io/ioutil"

	"google.golang.org/appengine/log"
)

// Server serves registered RPC services using registered codecs.
type Server struct {
	root     string
	services *serviceMap

	// ContextDecorator will be called as the last step of the creation of a new context.
	// If nil the context will not be decorated.
	ContextDecorator func(context.Context) (context.Context, error)
}

// NewServer returns a new RPC server.
func NewServer(root string) *Server {
	if root == "" {
		root = "/_ah/spi/"
	} else if root[len(root)-1] != '/' {
		root += "/"
	}

	server := &Server{root: root, services: new(serviceMap)}
	backend := newBackendService(server)
	server.services.register(backend, "BackendService", "", "", true, true)
	return server
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method has either 2 arguments and 2 return values:
//      *http.Request|Context, *arg => *reply, error
//      or 3 arguments and 1 return value:
//      *http.Request|Context, *arg, *reply => error
//    - The first argument is either *http.Request or Context.
//    - Second argument (*arg) and *reply are exported or local.
//    - First argument, *arg and *reply are all pointers.
//    - First (or second, if method has 2 arguments) return value is of type error.
//
// All other methods are ignored.
func (s *Server) RegisterService(srv interface{}, name, ver, desc string, isDefault bool) (*RPCService, error) {
	return s.services.register(srv, name, ver, desc, isDefault, false)
}

// RegisterServiceWithDefaults will register provided service and will try to
// infer Endpoints config params from its method names and types.
// See RegisterService for details.
func (s *Server) RegisterServiceWithDefaults(srv interface{}) (*RPCService, error) {
	return s.RegisterService(srv, "", "", "", true)
}

// Must is a helper that wraps a call to a function returning (*Template, error) and
// panics if the error is non-nil. It is intended for use in variable initializations
// such as:
// 	var s = endpoints.Must(endpoints.RegisterService(s, "Service", "v1", "some service", true))
//
func Must(s *RPCService, err error) *RPCService {
	if err != nil {
		panic(err)
	}
	return s
}

// ServiceByName returns a registered service or nil if there's no service
// registered by that name.
func (s *Server) ServiceByName(serviceName string) *RPCService {
	return s.services.serviceByName(serviceName)
}

// HandleHTTP adds Server s to specified http.ServeMux.
// If no mux is provided http.DefaultServeMux will be used.
func (s *Server) HandleHTTP(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}
	mux.Handle(s.root, s)
}

// ServeHTTP is Server's implementation of http.Handler interface.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	c := NewContext(r)
	if s.ContextDecorator != nil {
		ctx, err := s.ContextDecorator(c)
		if err != nil {
			writeError(w, err)
			return
		}
		c = ctx
	}

	// Always respond with JSON, even when an error occurs.
	// Note: API server doesn't expect an encoding in Content-Type header.
	w.Header().Set("Content-Type", "application/json")

	if r.Method != "POST" {
		err := fmt.Errorf("rpc: POST method required, got %q", r.Method)
		writeError(w, err)
		return
	}

	// methodName has "ServiceName.MethodName" format.
	var methodName string
	idx := strings.LastIndex(r.URL.Path, "/")
	if idx < 0 {
		writeError(w, fmt.Errorf("rpc: no method in path %q", r.URL.Path))
		return
	}
	methodName = r.URL.Path[idx+1:]

	// Get service method specs
	serviceSpec, methodSpec, err := s.services.get(methodName)
	if err != nil {
		writeError(w, err)
		return
	}

	// Initialize RPC method request
	reqValue := reflect.New(methodSpec.ReqType)

	body, err := ioutil.ReadAll(r.Body)
	r.Body.Close()
	if err != nil {
		writeError(w, err)
		return
	}
	log.Debugf(c, "SPI request body: %s", body)

	// if err := json.NewDecoder(r.Body).Decode(req.Interface()); err != nil {
	// 	writeError(w, fmt.Errorf("Error while decoding JSON: %q", err))
	// 	return
	// }
	if err := json.Unmarshal(body, reqValue.Interface()); err != nil {
		writeError(w, err)
		return
	}

	if err := validateRequest(reqValue.Interface()); err != nil {
		writeError(w, err)
		return
	}

	// Restore the body in the original request.
	r.Body = ioutil.NopCloser(bytes.NewReader(body))

	numIn, numOut := methodSpec.method.Type.NumIn(), methodSpec.method.Type.NumOut()
	// Construct arguments for the method call
	var httpReqOrCtx interface{} = r
	if methodSpec.wantsContext {
		httpReqOrCtx = c
	}
	args := []reflect.Value{serviceSpec.rcvr, reflect.ValueOf(httpReqOrCtx)}
	if numIn > 2 {
		args = append(args, reqValue)
	}

	var respValue reflect.Value
	if numIn > 3 {
		respValue = reflect.New(methodSpec.RespType)
		args = append(args, respValue)
	}

	// Invoke the service method
	var errValue reflect.Value
	res := methodSpec.method.Func.Call(args)
	if numOut == 2 {
		respValue = res[0]
		errValue = res[1]
	} else {
		errValue = res[0]
	}

	// Check if method returned an error
	if err := errValue.Interface(); err != nil {
		writeError(w, err.(error))
		return
	}

	// Encode non-error response
	if numIn == 4 || numOut == 2 {
		if err := json.NewEncoder(w).Encode(respValue.Interface()); err != nil {
			writeError(w, err)
		}
	}
}

// DefaultServer is the default RPC server, so you don't have to explicitly
// create one.
var DefaultServer *Server

// RegisterService registers a service using DefaultServer.
// See Server.RegisterService for details.
func RegisterService(srv interface{}, name, ver, desc string, isDefault bool) (
	*RPCService, error) {

	return DefaultServer.RegisterService(srv, name, ver, desc, isDefault)
}

// RegisterServiceWithDefaults registers a service using DefaultServer.
// See Server.RegisterServiceWithDefaults for details.
func RegisterServiceWithDefaults(srv interface{}) (*RPCService, error) {
	return DefaultServer.RegisterServiceWithDefaults(srv)
}

// HandleHTTP calls DefaultServer's HandleHTTP method using default serve mux.
func HandleHTTP() {
	DefaultServer.HandleHTTP(nil)
}

// TODO: var DefaultServer = NewServer("") won't work so it's in the init()
// function for now.
func init() {
	DefaultServer = NewServer("")
}
<commit_msg>always set content-type header before writing errors<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package endpoints

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"reflect"
	"strings"

	"golang.org/x/net/context"
	// Mainly for debug logging
	"io/ioutil"

	"google.golang.org/appengine/log"
)

// Server serves registered RPC services using registered codecs.
type Server struct {
	root     string
	services *serviceMap

	// ContextDecorator will be called as the last step of the creation of a new context.
	// If nil the context will not be decorated.
	ContextDecorator func(context.Context) (context.Context, error)
}

// NewServer returns a new RPC server.
func NewServer(root string) *Server {
	if root == "" {
		root = "/_ah/spi/"
	} else if root[len(root)-1] != '/' {
		root += "/"
	}

	server := &Server{root: root, services: new(serviceMap)}
	backend := newBackendService(server)
	server.services.register(backend, "BackendService", "", "", true, true)
	return server
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method has either 2 arguments and 2 return values:
//      *http.Request|Context, *arg => *reply, error
//      or 3 arguments and 1 return value:
//      *http.Request|Context, *arg, *reply => error
//    - The first argument is either *http.Request or Context.
//    - Second argument (*arg) and *reply are exported or local.
//    - First argument, *arg and *reply are all pointers.
//    - First (or second, if method has 2 arguments) return value is of type error.
//
// All other methods are ignored.
func (s *Server) RegisterService(srv interface{}, name, ver, desc string, isDefault bool) (*RPCService, error) {
	return s.services.register(srv, name, ver, desc, isDefault, false)
}

// RegisterServiceWithDefaults will register provided service and will try to
// infer Endpoints config params from its method names and types.
// See RegisterService for details.
func (s *Server) RegisterServiceWithDefaults(srv interface{}) (*RPCService, error) {
	return s.RegisterService(srv, "", "", "", true)
}

// Must is a helper that wraps a call to a function returning (*Template, error) and
// panics if the error is non-nil. It is intended for use in variable initializations
// such as:
// 	var s = endpoints.Must(endpoints.RegisterService(s, "Service", "v1", "some service", true))
//
func Must(s *RPCService, err error) *RPCService {
	if err != nil {
		panic(err)
	}
	return s
}

// ServiceByName returns a registered service or nil if there's no service
// registered by that name.
func (s *Server) ServiceByName(serviceName string) *RPCService {
	return s.services.serviceByName(serviceName)
}

// HandleHTTP adds Server s to specified http.ServeMux.
// If no mux is provided http.DefaultServeMux will be used.
func (s *Server) HandleHTTP(mux *http.ServeMux) {
	if mux == nil {
		mux = http.DefaultServeMux
	}
	mux.Handle(s.root, s)
}

// ServeHTTP is Server's implementation of http.Handler interface.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Always respond with JSON, even when an error occurs.
	// Note: API server doesn't expect an encoding in Content-Type header.
	w.Header().Set("Content-Type", "application/json")

	c := NewContext(r)
	if s.ContextDecorator != nil {
		ctx, err := s.ContextDecorator(c)
		if err != nil {
			writeError(w, err)
			return
		}
		c = ctx
	}

	if r.Method != "POST" {
		err := fmt.Errorf("rpc: POST method required, got %q", r.Method)
		writeError(w, err)
		return
	}

	// methodName has "ServiceName.MethodName" format.
	var methodName string
	idx := strings.LastIndex(r.URL.Path, "/")
	if idx < 0 {
		writeError(w, fmt.Errorf("rpc: no method in path %q", r.URL.Path))
		return
	}
	methodName = r.URL.Path[idx+1:]

	// Get service method specs
	serviceSpec, methodSpec, err := s.services.get(methodName)
	if err != nil {
		writeError(w, err)
		return
	}

	// Initialize RPC method request
	reqValue := reflect.New(methodSpec.ReqType)

	body, err := ioutil.ReadAll(r.Body)
	r.Body.Close()
	if err != nil {
		writeError(w, err)
		return
	}
	log.Debugf(c, "SPI request body: %s", body)

	// if err := json.NewDecoder(r.Body).Decode(req.Interface()); err != nil {
	// 	writeError(w, fmt.Errorf("Error while decoding JSON: %q", err))
	// 	return
	// }
	if err := json.Unmarshal(body, reqValue.Interface()); err != nil {
		writeError(w, err)
		return
	}

	if err := validateRequest(reqValue.Interface()); err != nil {
		writeError(w, err)
		return
	}

	// Restore the body in the original request.
	r.Body = ioutil.NopCloser(bytes.NewReader(body))

	numIn, numOut := methodSpec.method.Type.NumIn(), methodSpec.method.Type.NumOut()
	// Construct arguments for the method call
	var httpReqOrCtx interface{} = r
	if methodSpec.wantsContext {
		httpReqOrCtx = c
	}
	args := []reflect.Value{serviceSpec.rcvr, reflect.ValueOf(httpReqOrCtx)}
	if numIn > 2 {
		args = append(args, reqValue)
	}

	var respValue reflect.Value
	if numIn > 3 {
		respValue = reflect.New(methodSpec.RespType)
		args = append(args, respValue)
	}

	// Invoke the service method
	var errValue reflect.Value
	res := methodSpec.method.Func.Call(args)
	if numOut == 2 {
		respValue = res[0]
		errValue = res[1]
	} else {
		errValue = res[0]
	}

	// Check if method returned an error
	if err := errValue.Interface(); err != nil {
		writeError(w, err.(error))
		return
	}

	// Encode non-error response
	if numIn == 4 || numOut == 2 {
		if err := json.NewEncoder(w).Encode(respValue.Interface()); err != nil {
			writeError(w, err)
		}
	}
}

// DefaultServer is the default RPC server, so you don't have to explicitly
// create one.
var DefaultServer *Server

// RegisterService registers a service using DefaultServer.
// See Server.RegisterService for details.
func RegisterService(srv interface{}, name, ver, desc string, isDefault bool) (
	*RPCService, error) {

	return DefaultServer.RegisterService(srv, name, ver, desc, isDefault)
}

// RegisterServiceWithDefaults registers a service using DefaultServer.
// See Server.RegisterServiceWithDefaults for details.
func RegisterServiceWithDefaults(srv interface{}) (*RPCService, error) {
	return DefaultServer.RegisterServiceWithDefaults(srv)
}

// HandleHTTP calls DefaultServer's HandleHTTP method using default serve mux.
func HandleHTTP() {
	DefaultServer.HandleHTTP(nil)
}

// TODO: var DefaultServer = NewServer("") won't work so it's in the init()
// function for now.
func init() {
	DefaultServer = NewServer("")
}
<|endoftext|>


<commit_before>// Copyright 2022 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package migrations

import (
	"crypto/elliptic"
	"encoding/base32"
	"fmt"
	"strings"

	"code.gitea.io/gitea/modules/timeutil"

	"github.com/tstranex/u2f"
	"xorm.io/xorm"
	"xorm.io/xorm/schemas"
)

// v208 migration was completely broken
func remigrateU2FCredentials(x *xorm.Engine) error {
	// Create webauthnCredential table
	type webauthnCredential struct {
		ID              int64 `xorm:"pk autoincr"`
		Name            string
		LowerName       string `xorm:"unique(s)"`
		UserID          int64  `xorm:"INDEX unique(s)"`
		CredentialID    string `xorm:"INDEX VARCHAR(410)"` // CredentalID in U2F is at most 255bytes / 5 * 8 = 408 - add a few extra characters for safety
		PublicKey       []byte
		AttestationType string
		AAGUID          []byte
		SignCount       uint32 `xorm:"BIGINT"`
		CloneWarning    bool
		CreatedUnix     timeutil.TimeStamp `xorm:"INDEX created"`
		UpdatedUnix     timeutil.TimeStamp `xorm:"INDEX updated"`
	}
	if err := x.Sync2(&webauthnCredential{}); err != nil {
		return err
	}

	switch x.Dialect().URI().DBType {
	case schemas.MYSQL:
		_, err := x.Exec("ALTER TABLE webauthn_credential MODIFY COLUMN credential_id VARCHAR(410)")
		if err != nil {
			return err
		}
	case schemas.ORACLE:
		_, err := x.Exec("ALTER TABLE webauthn_credential MODIFY credential_id VARCHAR(410)")
		if err != nil {
			return err
		}
	case schemas.MSSQL:
		// This column has an index on it. I could write all of the code to attempt to change the index OR
		// I could just use recreate table.
		sess := x.NewSession()
		if err := sess.Begin(); err != nil {
			_ = sess.Close()
			return err
		}

		if err := recreateTable(sess, new(webauthnCredential)); err != nil {
			_ = sess.Close()
			return err
		}
		if err := sess.Commit(); err != nil {
			_ = sess.Close()
			return err
		}
		if err := sess.Close(); err != nil {
			return err
		}
	case schemas.POSTGRES:
		_, err := x.Exec("ALTER TABLE webauthn_credential ALTER COLUMN credential_id TYPE VARCHAR(410)")
		if err != nil {
			return err
		}
	default:
		// SQLite doesn't support ALTER COLUMN, and it already makes String _TEXT_ by default so no migration needed
		// nor is there any need to re-migrate
	}

	exist, err := x.IsTableExist("u2f_registration")
	if err != nil {
		return err
	}
	if !exist {
		return nil
	}

	// Now migrate the old u2f registrations to the new format
	type u2fRegistration struct {
		ID          int64 `xorm:"pk autoincr"`
		Name        string
		UserID      int64 `xorm:"INDEX"`
		Raw         []byte
		Counter     uint32             `xorm:"BIGINT"`
		CreatedUnix timeutil.TimeStamp `xorm:"INDEX created"`
		UpdatedUnix timeutil.TimeStamp `xorm:"INDEX updated"`
	}

	var start int
	regs := make([]*u2fRegistration, 0, 50)
	for {
		err := x.OrderBy("id").Limit(50, start).Find(&regs)
		if err != nil {
			return err
		}

		err = func() error {
			sess := x.NewSession()
			defer sess.Close()
			if err := sess.Begin(); err != nil {
				return fmt.Errorf("unable to allow start session. Error: %w", err)
			}
			if x.Dialect().URI().DBType == schemas.MSSQL {
				if _, err := sess.Exec("SET IDENTITY_INSERT `webauthn_credential` ON"); err != nil {
					return fmt.Errorf("unable to allow identity insert on webauthn_credential. Error: %w", err)
				}
			}
			for _, reg := range regs {
				parsed := new(u2f.Registration)
				err = parsed.UnmarshalBinary(reg.Raw)
				if err != nil {
					continue
				}
				remigrated := &webauthnCredential{
					ID:              reg.ID,
					Name:            reg.Name,
					LowerName:       strings.ToLower(reg.Name),
					UserID:          reg.UserID,
					CredentialID:    base32.HexEncoding.EncodeToString(parsed.KeyHandle),
					PublicKey:       elliptic.Marshal(elliptic.P256(), parsed.PubKey.X, parsed.PubKey.Y),
					AttestationType: "fido-u2f",
					AAGUID:          []byte{},
					SignCount:       reg.Counter,
					UpdatedUnix:     reg.UpdatedUnix,
					CreatedUnix:     reg.CreatedUnix,
				}

				has, err := sess.ID(reg.ID).Where("id = ?", reg.ID).Get(new(webauthnCredential))
				if err != nil {
					return fmt.Errorf("unable to get webauthn_credential[%d]. Error: %w", reg.ID, err)
				}
				if !has {
					_, err = sess.Insert(remigrated)
					if err != nil {
						return fmt.Errorf("unable to (re)insert webauthn_credential[%d]. Error: %w", reg.ID, err)
					}

					continue
				}

				_, err = sess.ID(remigrated.ID).AllCols().Update(remigrated)
				if err != nil {
					return fmt.Errorf("unable to update webauthn_credential[%d]. Error: %w", reg.ID, err)
				}
			}
			return sess.Commit()
		}()
		if err != nil {
			return err
		}

		if len(regs) < 50 {
			break
		}
		start += 50
		regs = regs[:0]
	}

	return nil
}
<commit_msg>Fix migration v210 (#18892)<commit_after>// Copyright 2022 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package migrations

import (
	"crypto/elliptic"
	"encoding/base32"
	"fmt"
	"strings"

	"code.gitea.io/gitea/modules/timeutil"

	"github.com/tstranex/u2f"
	"xorm.io/xorm"
	"xorm.io/xorm/schemas"
)

// v208 migration was completely broken
func remigrateU2FCredentials(x *xorm.Engine) error {
	// Create webauthnCredential table
	type webauthnCredential struct {
		ID              int64 `xorm:"pk autoincr"`
		Name            string
		LowerName       string `xorm:"unique(s)"`
		UserID          int64  `xorm:"INDEX unique(s)"`
		CredentialID    string `xorm:"INDEX VARCHAR(410)"` // CredentalID in U2F is at most 255bytes / 5 * 8 = 408 - add a few extra characters for safety
		PublicKey       []byte
		AttestationType string
		AAGUID          []byte
		SignCount       uint32 `xorm:"BIGINT"`
		CloneWarning    bool
		CreatedUnix     timeutil.TimeStamp `xorm:"INDEX created"`
		UpdatedUnix     timeutil.TimeStamp `xorm:"INDEX updated"`
	}
	if err := x.Sync2(&webauthnCredential{}); err != nil {
		return err
	}

	switch x.Dialect().URI().DBType {
	case schemas.MYSQL:
		_, err := x.Exec("ALTER TABLE webauthn_credential MODIFY COLUMN credential_id VARCHAR(410)")
		if err != nil {
			return err
		}
	case schemas.ORACLE:
		_, err := x.Exec("ALTER TABLE webauthn_credential MODIFY credential_id VARCHAR(410)")
		if err != nil {
			return err
		}
	case schemas.MSSQL:
		// This column has an index on it. I could write all of the code to attempt to change the index OR
		// I could just use recreate table.
		sess := x.NewSession()
		if err := sess.Begin(); err != nil {
			_ = sess.Close()
			return err
		}

		if err := recreateTable(sess, new(webauthnCredential)); err != nil {
			_ = sess.Close()
			return err
		}
		if err := sess.Commit(); err != nil {
			_ = sess.Close()
			return err
		}
		if err := sess.Close(); err != nil {
			return err
		}
	case schemas.POSTGRES:
		_, err := x.Exec("ALTER TABLE webauthn_credential ALTER COLUMN credential_id TYPE VARCHAR(410)")
		if err != nil {
			return err
		}
	default:
		// SQLite doesn't support ALTER COLUMN, and it already makes String _TEXT_ by default so no migration needed
		// nor is there any need to re-migrate
	}

	exist, err := x.IsTableExist("u2f_registration")
	if err != nil {
		return err
	}
	if !exist {
		return nil
	}

	// Now migrate the old u2f registrations to the new format
	type u2fRegistration struct {
		ID          int64 `xorm:"pk autoincr"`
		Name        string
		UserID      int64 `xorm:"INDEX"`
		Raw         []byte
		Counter     uint32             `xorm:"BIGINT"`
		CreatedUnix timeutil.TimeStamp `xorm:"INDEX created"`
		UpdatedUnix timeutil.TimeStamp `xorm:"INDEX updated"`
	}

	var start int
	regs := make([]*u2fRegistration, 0, 50)
	for {
		err := x.OrderBy("id").Limit(50, start).Find(&regs)
		if err != nil {
			return err
		}

		err = func() error {
			sess := x.NewSession()
			defer sess.Close()
			if err := sess.Begin(); err != nil {
				return fmt.Errorf("unable to allow start session. Error: %w", err)
			}
			if x.Dialect().URI().DBType == schemas.MSSQL {
				if _, err := sess.Exec("SET IDENTITY_INSERT `webauthn_credential` ON"); err != nil {
					return fmt.Errorf("unable to allow identity insert on webauthn_credential. Error: %w", err)
				}
			}
			for _, reg := range regs {
				parsed := new(u2f.Registration)
				err = parsed.UnmarshalBinary(reg.Raw)
				if err != nil {
					continue
				}
				remigrated := &webauthnCredential{
					ID:              reg.ID,
					Name:            reg.Name,
					LowerName:       strings.ToLower(reg.Name),
					UserID:          reg.UserID,
					CredentialID:    base32.HexEncoding.EncodeToString(parsed.KeyHandle),
					PublicKey:       elliptic.Marshal(elliptic.P256(), parsed.PubKey.X, parsed.PubKey.Y),
					AttestationType: "fido-u2f",
					AAGUID:          []byte{},
					SignCount:       reg.Counter,
					UpdatedUnix:     reg.UpdatedUnix,
					CreatedUnix:     reg.CreatedUnix,
				}

				has, err := sess.ID(reg.ID).Get(new(webauthnCredential))
				if err != nil {
					return fmt.Errorf("unable to get webauthn_credential[%d]. Error: %w", reg.ID, err)
				}
				if !has {
					has, err := sess.Where("`lower_name`=?", remigrated.LowerName).And("`user_id`=?", remigrated.UserID).Exist(new(webauthnCredential))
					if err != nil {
						return fmt.Errorf("unable to check webauthn_credential[lower_name: %s, user_id:%v]. Error: %w", remigrated.LowerName, remigrated.UserID, err)
					}
					if !has {
						_, err = sess.Insert(remigrated)
						if err != nil {
							return fmt.Errorf("unable to (re)insert webauthn_credential[%d]. Error: %w", reg.ID, err)
						}

						continue
					}
				}

				_, err = sess.ID(remigrated.ID).AllCols().Update(remigrated)
				if err != nil {
					return fmt.Errorf("unable to update webauthn_credential[%d]. Error: %w", reg.ID, err)
				}
			}
			return sess.Commit()
		}()
		if err != nil {
			return err
		}

		if len(regs) < 50 {
			break
		}
		start += 50
		regs = regs[:0]
	}

	return nil
}
<|endoftext|>


<commit_before>// +build !windows

// Functions shared between linux/darwin.
package allocdir

import (
	"fmt"
	"os"
	"os/user"
	"strconv"
	"syscall"
)

func (d *AllocDir) linkOrCopy(src, dst string, perm os.FileMode) error {
	// Attempt to hardlink.
	if err := os.Link(src, dst); err == nil {
		return nil
	}

	return fileCopy(src, dst, perm)
}

func (d *AllocDir) dropDirPermissions(path string) error {
	// Can't do anything if not root.
	if syscall.Geteuid() != 0 {
		return nil
	}

	u, err := user.Lookup("nobody")
	if err != nil {
		return err
	}

	uid, err := getUid(u)
	if err != nil {
		return err
	}

	gid, err := getGid(u)
	if err != nil {
		return err
	}

	if err := os.Chown(path, uid, gid); err != nil {
		return fmt.Errorf("Couldn't change owner/group of %v to (uid: %v, gid: %v): %v", path, uid, gid, err)
	}

	if err := os.Chmod(path, 0777); err != nil {
		return fmt.Errorf("Chmod(%v) failed: %v", path, err)
	}

	return nil
}

func getUid(u *user.User) (int, error) {
	uid, err := strconv.Atoi(u.Uid)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert Uid to an int: %v", err)
	}

	return uid, nil
}

func getGid(u *user.User) (int, error) {
	gid, err := strconv.Atoi(u.Gid)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert Gid to an int: %v", err)
	}

	return gid, nil
}
<commit_msg>Remove CGO dependency for user lookup in allocdir<commit_after>// +build !windows

// Functions shared between linux/darwin.
package allocdir

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"strconv"
	"strings"
	"syscall"
)

func (d *AllocDir) linkOrCopy(src, dst string, perm os.FileMode) error {
	// Attempt to hardlink.
	if err := os.Link(src, dst); err == nil {
		return nil
	}

	return fileCopy(src, dst, perm)
}

func (d *AllocDir) dropDirPermissions(path string) error {
	// Can't do anything if not root.
	if syscall.Geteuid() != 0 {
		return nil
	}

	u, err := userLookup("nobody")
	if err != nil {
		return err
	}

	uid, err := getUid(u)
	if err != nil {
		return err
	}

	gid, err := getGid(u)
	if err != nil {
		return err
	}

	if err := os.Chown(path, uid, gid); err != nil {
		return fmt.Errorf("Couldn't change owner/group of %v to (uid: %v, gid: %v): %v", path, uid, gid, err)
	}

	if err := os.Chmod(path, 0777); err != nil {
		return fmt.Errorf("Chmod(%v) failed: %v", path, err)
	}

	return nil
}

func getUid(u *user.User) (int, error) {
	uid, err := strconv.Atoi(u.Uid)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert Uid to an int: %v", err)
	}

	return uid, nil
}

func getGid(u *user.User) (int, error) {
	gid, err := strconv.Atoi(u.Gid)
	if err != nil {
		return 0, fmt.Errorf("Unable to convert Gid to an int: %v", err)
	}

	return gid, nil
}

// userLookup checks if the given username or uid is present in /etc/passwd
// and returns the user struct.
// If the username is not found, an error is returned.
// Credit to @creak, https://github.com/docker/docker/pull/1096
func userLookup(uid string) (*user.User, error) {
	file, err := ioutil.ReadFile("/etc/passwd")
	if err != nil {
		return nil, err
	}

	for _, line := range strings.Split(string(file), "\n") {
		data := strings.Split(line, ":")
		if len(data) > 5 && (data[0] == uid || data[2] == uid) {
			return &user.User{
				Uid:      data[2],
				Gid:      data[3],
				Username: data[0],
				Name:     data[4],
				HomeDir:  data[5],
			}, nil
		}
	}
	return nil, fmt.Errorf("User not found in /etc/passwd")
}
<|endoftext|>


<commit_before>package core

// Arguments represents a structured set of arguments passed to a predicate.
type Arguments struct {
	positionals   []*Thunk
	expandedList  *Thunk
	keywords      []KeywordArgument
	expandedDicts []*Thunk
}

// NewArguments creates a new Arguments.
func NewArguments(
	ps []PositionalArgument,
	ks []KeywordArgument,
	ds []*Thunk) Arguments {
	ts := make([]*Thunk, 0, len(ps))
	l := (*Thunk)(nil)

	for i, p := range ps {
		if p.expanded {
			l = mergeRestPositionalArgs(ps[i].value, ps[i+1:]...)
			break
		}

		ts = append(ts, p.value)
	}

	return Arguments{ts, l, ks, ds}
}

func mergeRestPositionalArgs(t *Thunk, ps ...PositionalArgument) *Thunk {
	for _, p := range ps {
		if p.expanded {
			t = PApp(Merge, t, p.value)
		} else {
			t = PApp(
				NewLazyFunction(appendFuncSignature, appendFunc), // Avoid initialization loop
				t, p.value)
		}
	}

	return t
}

func (args *Arguments) nextPositional() *Thunk {
	if len(args.positionals) != 0 {
		defer func() { args.positionals = args.positionals[1:] }()
		return args.positionals[0]
	}

	if args.expandedList == nil {
		return nil
	}

	l := args.expandedList
	args.expandedList = PApp(Rest, l)
	return PApp(First, l)
}

func (args *Arguments) restPositionals() *Thunk {
	ps := args.positionals
	l := args.expandedList
	args.positionals = nil
	args.expandedList = nil

	if l == nil {
		return NewList(ps...)
	}

	return PApp(Merge, NewList(ps...), l)
}

func (args *Arguments) searchKeyword(s string) *Thunk {
	for i, k := range args.keywords {
		if s == k.name {
			args.keywords = append(args.keywords[:i], args.keywords[i+1:]...)
			return k.value
		}
	}

	for i, t := range args.expandedDicts {
		v := t.Eval()
		d, ok := v.(DictionaryType)

		if !ok {
			return NotDictionaryError(v)
		}

		k := StringType(s)

		if v, ok := d.Search(k); ok {
			new := make([]*Thunk, len(args.expandedDicts))
			copy(new, args.expandedDicts)
			new[i] = Normal(d.Remove(k))
			args.expandedDicts = new
			return v
		}
	}

	return nil
}

func (args *Arguments) restKeywords() *Thunk {
	defer func() {
		args.keywords = nil
		args.expandedDicts = nil
	}()

	t := EmptyDictionary

	for _, k := range args.keywords {
		t = PApp(Insert, t, NewString(k.name), k.value)
	}

	for _, tt := range args.expandedDicts {
		t = PApp(Merge, t, tt)
	}

	return t
}

// Merge merges 2 sets of arguments into one.
func (args Arguments) Merge(merged Arguments) Arguments {
	var new Arguments

	if new.expandedList == nil {
		new.positionals = append(args.positionals, merged.positionals...)
		new.expandedList = merged.expandedList
	} else {
		new.positionals = args.positionals
		new.expandedList = PApp(
			Append,
			append([]*Thunk{args.expandedList}, merged.positionals...)...)

		if merged.expandedList != nil {
			new.expandedList = PApp(Merge, new.expandedList, merged.expandedList)
		}
	}

	new.keywords = append(args.keywords, merged.keywords...)
	new.expandedDicts = append(args.expandedDicts, merged.expandedDicts...)

	return new
}

func (args Arguments) empty() *Thunk {
	if args.positionals != nil && len(args.positionals) > 0 {
		return argumentError("%d positional arguments are left", len(args.positionals))
	}

	// Testing args.expandedList is impossible because we cannot know its length
	// without evaluating it.

	n := 0

	if args.expandedDicts != nil {
		for _, t := range args.expandedDicts {
			v := t.Eval()
			d, ok := v.(DictionaryType)

			if !ok {
				return NotDictionaryError(v)
			}

			n += d.Size()
		}
	}

	if n != 0 || args.keywords != nil && len(args.keywords) > 0 {
		return argumentError("%d keyword arguments are left", len(args.keywords)+n)
	}

	return nil
}
<commit_msg>Refactor arguments.go<commit_after>package core

// Arguments represents a structured set of arguments passed to a predicate.
type Arguments struct {
	positionals   []*Thunk
	expandedList  *Thunk
	keywords      []KeywordArgument
	expandedDicts []*Thunk
}

// NewArguments creates a new Arguments.
func NewArguments(
	ps []PositionalArgument,
	ks []KeywordArgument,
	ds []*Thunk) Arguments {
	ts := make([]*Thunk, 0, len(ps))
	l := (*Thunk)(nil)

	for i, p := range ps {
		if p.expanded {
			l = mergeRestPositionalArgs(ps[i].value, ps[i+1:]...)
			break
		}

		ts = append(ts, p.value)
	}

	return Arguments{ts, l, ks, ds}
}

func mergeRestPositionalArgs(t *Thunk, ps ...PositionalArgument) *Thunk {
	for _, p := range ps {
		if p.expanded {
			t = PApp(Merge, t, p.value)
		} else {
			t = PApp(
				NewLazyFunction(appendFuncSignature, appendFunc), // Avoid initialization loop
				t, p.value)
		}
	}

	return t
}

func (args *Arguments) nextPositional() *Thunk {
	if len(args.positionals) != 0 {
		t := args.positionals[0]
		args.positionals = args.positionals[1:]
		return t
	}

	if args.expandedList == nil {
		return nil
	}

	l := args.expandedList
	args.expandedList = PApp(Rest, l)
	return PApp(First, l)
}

func (args *Arguments) restPositionals() *Thunk {
	ts := args.positionals
	l := args.expandedList
	args.positionals = nil
	args.expandedList = nil

	if l == nil {
		return NewList(ts...)
	}

	return PApp(Merge, NewList(ts...), l)
}

func (args *Arguments) searchKeyword(s string) *Thunk {
	for i, k := range args.keywords {
		if s == k.name {
			args.keywords = append(args.keywords[:i], args.keywords[i+1:]...)
			return k.value
		}
	}

	for i, t := range args.expandedDicts {
		v := t.Eval()
		d, ok := v.(DictionaryType)

		if !ok {
			return NotDictionaryError(v)
		}

		k := StringType(s)

		if v, ok := d.Search(k); ok {
			ds := make([]*Thunk, len(args.expandedDicts))
			copy(ds, args.expandedDicts)
			ds[i] = Normal(d.Remove(k))
			args.expandedDicts = ds
			return v
		}
	}

	return nil
}

func (args *Arguments) restKeywords() *Thunk {
	ks := args.keywords
	ds := args.expandedDicts
	args.keywords = nil
	args.expandedDicts = nil

	t := EmptyDictionary

	for _, k := range ks {
		t = PApp(Insert, t, NewString(k.name), k.value)
	}

	for _, d := range ds {
		t = PApp(Merge, t, d)
	}

	return t
}

// Merge merges 2 sets of arguments into one.
func (args Arguments) Merge(old Arguments) Arguments {
	var new Arguments

	if new.expandedList == nil {
		new.positionals = append(args.positionals, old.positionals...)
		new.expandedList = old.expandedList
	} else {
		new.positionals = args.positionals
		new.expandedList = PApp(
			Append,
			append([]*Thunk{args.expandedList}, old.positionals...)...)

		if old.expandedList != nil {
			new.expandedList = PApp(Merge, new.expandedList, old.expandedList)
		}
	}

	new.keywords = append(args.keywords, old.keywords...)
	new.expandedDicts = append(args.expandedDicts, old.expandedDicts...)

	return new
}

func (args Arguments) empty() *Thunk {
	if len(args.positionals) > 0 {
		return argumentError("%d positional arguments are left", len(args.positionals))
	}

	// Testing args.expandedList is impossible because we cannot know its length
	// without evaluating it.

	n := 0

	for _, t := range args.expandedDicts {
		v := t.Eval()
		d, ok := v.(DictionaryType)

		if !ok {
			return NotDictionaryError(v)
		}

		n += d.Size()
	}

	if n != 0 || args.keywords != nil && len(args.keywords) > 0 {
		return argumentError("%d keyword arguments are left", len(args.keywords)+n)
	}

	return nil
}
<|endoftext|>


<commit_before>package config

import "testing"

type fakeConfig struct {
	FieldOne string `servconf:"field_one,required"`
	FieldTwo string `servconf:"field_two"`
}

func TestServiceConfig(t *testing.T) {
	var confValidSC1 ServiceConfig
	var confValidSC2 ServiceConfig
	var confInvalidSC1 ServiceConfig
	confValidSC1 = map[string]interface{}{
		"field_one": "test1",
		"field_two": "test2",
	}
	confValidSC2 = map[string]interface{}{
		"field_one": "test1",
	}
	confInvalidSC1 = map[string]interface{}{
		"field_two": "test",
	}

	// TODO add other config structs with invalid tags and test it
	t.Run("AssignValidSC1", func(t *testing.T) {
		cfg := fakeConfig{}
		err := confValidSC1.Assign(&cfg)
		if err != nil {
			t.Errorf("Assign returned an unexpected error: %s", err)
		}
		if cfg.FieldOne != "test1" {
			t.Errorf("FieldOne value %q is not \"test1\"", cfg.FieldOne)
		}
		if cfg.FieldTwo != "test2" {
			t.Errorf("FieldTwo value %q is not \"test2\"", cfg.FieldTwo)
		}
	})
	t.Run("AssignValidSC2", func(t *testing.T) {
		cfg := fakeConfig{}
		err := confValidSC2.Assign(&cfg)
		if err != nil {
			t.Errorf("Assign returned an unexpected error: %s", err)
		}
		if cfg.FieldOne != "test1" {
			t.Errorf("FieldOne value %q is not \"test1\"", cfg.FieldOne)
		}
		if cfg.FieldTwo != "" {
			t.Errorf("FieldTwo value %q is not \"\"", cfg.FieldTwo)
		}
	})
	t.Run("AssignInvalidSC1", func(t *testing.T) {
		cfg := fakeConfig{}
		err := confInvalidSC1.Assign(&cfg)
		if err == nil {
			t.Errorf("Assign did not return an error with invalid input")
			return
		}
		switch err.(type) {
		case *InvalidTagError:
			if err.(*InvalidTagError).NameTag == "field_one" {
				return
			}
		}
		t.Errorf("Assign returned an unexpected error: %s", err)
	})
}
<commit_msg>Adjust ServiceConfig test to have a clearer error<commit_after>package config

import "testing"

type fakeConfig struct {
	FieldOne string `servconf:"field_one,required"`
	FieldTwo string `servconf:"field_two"`
}

func TestServiceConfig(t *testing.T) {
	var confValidSC1 ServiceConfig
	var confValidSC2 ServiceConfig
	var confInvalidSC1 ServiceConfig
	confValidSC1 = map[string]interface{}{
		"field_one": "test1",
		"field_two": "test2",
	}
	confValidSC2 = map[string]interface{}{
		"field_one": "test1",
	}
	confInvalidSC1 = map[string]interface{}{
		"field_two": "test",
	}

	// TODO add other config structs with invalid tags and test it
	t.Run("AssignValidSC1", func(t *testing.T) {
		cfg := fakeConfig{}
		err := confValidSC1.Assign(&cfg)
		if err != nil {
			t.Errorf("Assign returned an unexpected error: %s", err)
		}
		if cfg.FieldOne != "test1" {
			t.Errorf("FieldOne value %q is not \"test1\"", cfg.FieldOne)
		}
		if cfg.FieldTwo != "test2" {
			t.Errorf("FieldTwo value %q is not \"test2\"", cfg.FieldTwo)
		}
	})
	t.Run("AssignValidSC2", func(t *testing.T) {
		cfg := fakeConfig{}
		err := confValidSC2.Assign(&cfg)
		if err != nil {
			t.Errorf("Assign returned an unexpected error: %s", err)
		}
		if cfg.FieldOne != "test1" {
			t.Errorf("FieldOne value %q is not \"test1\"", cfg.FieldOne)
		}
		if cfg.FieldTwo != "" {
			t.Errorf("FieldTwo value %q is not \"\"", cfg.FieldTwo)
		}
	})
	t.Run("AssignInvalidSC1", func(t *testing.T) {
		cfg := fakeConfig{}
		err := confInvalidSC1.Assign(&cfg)
		if err == nil {
			t.Errorf("Assign did not return an error when missing a required field")
			return
		}
		switch err.(type) {
		case *InvalidTagError:
			if err.(*InvalidTagError).NameTag == "field_one" {
				return
			}
		}
		t.Errorf("Assign returned an unexpected error: %s", err)
	})
}
<|endoftext|>


<commit_before>package schema

import (
	"database/sql/driver"
	"fmt"

	"golang.org/x/net/context"

	"github.com/araddon/qlbridge/expr"
	"github.com/araddon/qlbridge/rel"
	"github.com/araddon/qlbridge/value"
)

var (
	ErrNotFound       = fmt.Errorf("Not Found")
	ErrNotImplemented = fmt.Errorf("Not Implemented")
)

type (
	// Source is factory registered to create connections to a backing dastasource.
	// Datasources are most likely a database, file, api, in-mem data etc.
	// It is thread-safe, singleton, responsible for creating connections and
	// exposing schema and performing ddl operations. It also exposes partition information
	// optionally if a distributed source, or able to be distributed queries.
	//
	// DDL/Schema Operations
	// - schema discovery, tables, columns etc
	// - create
	// - index
	//
	// Lifecycle:
	// - Init()
	// - Setup()
	// - running ....
	// - Close()
	Source interface {
		// Init provides opportunity for those sources that require
		// no configuration and sniff schema from their environment time
		// to load pre-schema discovery
		Init()
		// Setup A Datasource optional interface for getting the SourceSchema injected
		// during creation/starup.  Since the Source is a singleton, stateful manager
		// it has a startup/shutdown process.
		Setup(*SchemaSource) error
		// Close
		Close() error
		// Open create a connection (not thread safe) to this source
		Open(source string) (Conn, error)
		// List of tables provided by this source
		Tables() []string
		// provides table schema info
		Table(table string) (*Table, error)
		// Create/Alter TODO
	}
	// SourcePartitionable DataSource that is partitionable into ranges for splitting
	//  reads, writes onto different nodes.
	SourcePartitionable interface {
		// Many databases's already have internal Partitions, allow those to
		// be exposed for use in our partitioning
		Partitions() []*Partition
		PartitionSource(p *Partition) (Conn, error)
	}
	// SourceTableColumn is a partial source that just provides access to
	// Column schema info, used in Generators.
	SourceTableColumn interface {
		// Underlying data type of column
		Column(col string) (value.ValueType, bool)
	}
)

type (
	// Conn A Connection/Session to a file, api, backend database.  Provides DML operations.
	//
	// Minimum Read Features to provide Sql Select:
	//  - Scanning:   iterate through messages/rows
	//  - Schema Tables:  at a minium list of tables available, the column level data
	//                    can be introspected so is optional
	//
	// Planning:
	//  - CreateMutator(ctx *plan.Context) :  execute a mutation task insert, delete, update
	//
	//
	// Non Select based Sql DML Operations for Mutator:
	//  - Deletion:    (sql delete)
	//      Delete()
	//      DeleteExpression()
	//  - Upsert Interface   (sql Update, Upsert, Insert)
	//      Put()
	//      PutMulti()
	//
	// DataSource Connection Session that is Stateful.  this is
	// really a generic interface, will actually implement features
	// below:  SchemaColumns, Scanner, Seeker, Mutator
	Conn interface {
		Close() error
	}
	// ConnAll interface
	ConnAll interface {
		Close() error
		ConnColumns
		Iterator
		ConnSeeker
		ConnUpsert
		ConnDeletion
	}
	// ConnColumns Interface for a data source connection exposing column positions for []driver.Value iteration
	ConnColumns interface {
		//Conn
		Columns() []string
	}
	// ConnScanner is the most basic of data sources, just iterate through
	//  rows without any optimizations.  Key-Value store like csv, redis, cassandra.
	ConnScanner interface {
		Conn
		Iterator
	}
	// ConnScannerIterator Another advanced iterator, probably deprecate?
	ConnScannerIterator interface {
		//Conn
		// create a new iterator to scan through row by row
		CreateIterator() Iterator
		MesgChan() <-chan Message
	}
	// ConnSeeker is a datsource that is Key-Value store, allows relational
	//  implementation to be faster for Seeking row values instead of scanning
	ConnSeeker interface {
		//Conn
		// Just because we have Get, Multi-Get, doesn't mean we can seek all
		// expressions, find out with CanSeek for given expression
		CanSeek(*rel.SqlSelect) bool
		Get(key driver.Value) (Message, error)
		MultiGet(keys []driver.Value) ([]Message, error)
	}
	// ConnMutation creates a Mutator connection similar to Open() connection for select
	//  - accepts the plan context used in this upsert/insert/update
	//  - returns a connection which must be closed
	ConnMutation interface {
		CreateMutator(pc interface{} /*plan.Context*/) (ConnMutator, error)
	}
	// ConnMutator Mutator Connection
	ConnMutator interface {
		ConnUpsert
		ConnDeletion
	}
	// ConnUpsert Mutation interface for Put
	//  - assumes datasource understands key(s?)
	ConnUpsert interface {
		Put(ctx context.Context, key Key, value interface{}) (Key, error)
		PutMulti(ctx context.Context, keys []Key, src interface{}) ([]Key, error)
	}
	// ConnPatchWhere pass through where expression to underlying datasource
	//  Used for update statements WHERE x = y
	ConnPatchWhere interface {
		PatchWhere(ctx context.Context, where expr.Node, patch interface{}) (int64, error)
	}
	// ConnDeletion deletion interface for data sources
	ConnDeletion interface {
		// Delete using this key
		Delete(driver.Value) (int, error)
		// Delete with given expression
		DeleteExpression(p interface{} /* plan.Delete */, n expr.Node) (int, error)
	}
)
<commit_msg>adding missing interface back in (#169)<commit_after>package schema

import (
	"database/sql/driver"
	"fmt"

	"golang.org/x/net/context"

	"github.com/araddon/qlbridge/expr"
	"github.com/araddon/qlbridge/rel"
	"github.com/araddon/qlbridge/value"
)

var (
	ErrNotFound       = fmt.Errorf("Not Found")
	ErrNotImplemented = fmt.Errorf("Not Implemented")
)

type (
	// Source is factory registered to create connections to a backing dastasource.
	// Datasources are most likely a database, file, api, in-mem data etc.
	// It is thread-safe, singleton, responsible for creating connections and
	// exposing schema and performing ddl operations. It also exposes partition information
	// optionally if a distributed source, or able to be distributed queries.
	//
	// DDL/Schema Operations
	// - schema discovery, tables, columns etc
	// - create
	// - index
	//
	// Lifecycle:
	// - Init()
	// - Setup()
	// - running ....
	// - Close()
	Source interface {
		// Init provides opportunity for those sources that require
		// no configuration and sniff schema from their environment time
		// to load pre-schema discovery
		Init()
		// Setup A Datasource optional interface for getting the SourceSchema injected
		// during creation/starup.  Since the Source is a singleton, stateful manager
		// it has a startup/shutdown process.
		Setup(*SchemaSource) error
		// Close
		Close() error
		// Open create a connection (not thread safe) to this source
		Open(source string) (Conn, error)
		// List of tables provided by this source
		Tables() []string
		// provides table schema info
		Table(table string) (*Table, error)
		// Create/Alter TODO
	}
	// SourceTableSchema Partial interface for just Table()
	SourceTableSchema interface {
		Table(table string) (*Table, error)
	}
	// SourcePartitionable DataSource that is partitionable into ranges for splitting
	//  reads, writes onto different nodes.
	SourcePartitionable interface {
		// Many databases's already have internal Partitions, allow those to
		// be exposed for use in our partitioning
		Partitions() []*Partition
		PartitionSource(p *Partition) (Conn, error)
	}
	// SourceTableColumn is a partial source that just provides access to
	// Column schema info, used in Generators.
	SourceTableColumn interface {
		// Underlying data type of column
		Column(col string) (value.ValueType, bool)
	}
)

type (
	// Conn A Connection/Session to a file, api, backend database.  Provides DML operations.
	//
	// Minimum Read Features to provide Sql Select:
	//  - Scanning:   iterate through messages/rows
	//  - Schema Tables:  at a minium list of tables available, the column level data
	//                    can be introspected so is optional
	//
	// Planning:
	//  - CreateMutator(ctx *plan.Context) :  execute a mutation task insert, delete, update
	//
	//
	// Non Select based Sql DML Operations for Mutator:
	//  - Deletion:    (sql delete)
	//      Delete()
	//      DeleteExpression()
	//  - Upsert Interface   (sql Update, Upsert, Insert)
	//      Put()
	//      PutMulti()
	//
	// DataSource Connection Session that is Stateful.  this is
	// really a generic interface, will actually implement features
	// below:  SchemaColumns, Scanner, Seeker, Mutator
	Conn interface {
		Close() error
	}
	// ConnAll interface
	ConnAll interface {
		Close() error
		ConnColumns
		Iterator
		ConnSeeker
		ConnUpsert
		ConnDeletion
	}
	// ConnColumns Interface for a data source connection exposing column positions for []driver.Value iteration
	ConnColumns interface {
		//Conn
		Columns() []string
	}
	// ConnScanner is the most basic of data sources, just iterate through
	//  rows without any optimizations.  Key-Value store like csv, redis, cassandra.
	ConnScanner interface {
		Conn
		Iterator
	}
	// ConnScannerIterator Another advanced iterator, probably deprecate?
	ConnScannerIterator interface {
		//Conn
		// create a new iterator to scan through row by row
		CreateIterator() Iterator
		MesgChan() <-chan Message
	}
	// ConnSeeker is a datsource that is Key-Value store, allows relational
	//  implementation to be faster for Seeking row values instead of scanning
	ConnSeeker interface {
		//Conn
		// Just because we have Get, Multi-Get, doesn't mean we can seek all
		// expressions, find out with CanSeek for given expression
		CanSeek(*rel.SqlSelect) bool
		Get(key driver.Value) (Message, error)
		MultiGet(keys []driver.Value) ([]Message, error)
	}
	// ConnMutation creates a Mutator connection similar to Open() connection for select
	//  - accepts the plan context used in this upsert/insert/update
	//  - returns a connection which must be closed
	ConnMutation interface {
		CreateMutator(pc interface{} /*plan.Context*/) (ConnMutator, error)
	}
	// ConnMutator Mutator Connection
	ConnMutator interface {
		ConnUpsert
		ConnDeletion
	}
	// ConnUpsert Mutation interface for Put
	//  - assumes datasource understands key(s?)
	ConnUpsert interface {
		Put(ctx context.Context, key Key, value interface{}) (Key, error)
		PutMulti(ctx context.Context, keys []Key, src interface{}) ([]Key, error)
	}
	// ConnPatchWhere pass through where expression to underlying datasource
	//  Used for update statements WHERE x = y
	ConnPatchWhere interface {
		PatchWhere(ctx context.Context, where expr.Node, patch interface{}) (int64, error)
	}
	// ConnDeletion deletion interface for data sources
	ConnDeletion interface {
		// Delete using this key
		Delete(driver.Value) (int, error)
		// Delete with given expression
		DeleteExpression(p interface{} /* plan.Delete */, n expr.Node) (int, error)
	}
)
<|endoftext|>


<commit_before>package syslog

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/segmentio/ecs-logs/lib"
)

const (
	DefaultTemplate = "<{{.PRIVAL}}>{{.TIMESTAMP}} {{.GROUP}}[{{.STREAM}}]: {{.MSG}}"
)

type WriterConfig struct {
	Network    string
	Address    string
	Template   string
	TimeFormat string
	TLS        *tls.Config
}

func NewWriter(group string, stream string) (w ecslogs.Writer, err error) {
	var c WriterConfig
	var s string
	var u *url.URL

	if s = os.Getenv("SYSLOG_URL"); len(s) != 0 {
		if u, err = url.Parse(s); err != nil {
			err = fmt.Errorf("invalid syslog URL: %s", err)
			return
		}

		c.Network = u.Scheme
		c.Address = u.Host
	}

	c.Template = os.Getenv("SYSLOG_TEMPLATE")
	c.TimeFormat = os.Getenv("SYSLOG_TIME_FORMAT")

	return DialWriter(c)
}

func DialWriter(config WriterConfig) (w ecslogs.Writer, err error) {
	var netopts []string
	var addropts []string
	var backend io.Writer

	if len(config.Address) != 0 {
		if len(config.Network) != 0 {
			netopts = []string{config.Network}
		} else {
			// When starting with a '/' we assume it's gonna be a file path,
			// otherwise we fallback to trying a TLS connection so we don't
			// implicitly send logs over an unsecured link.
			if strings.HasPrefix(config.Address, "/") {
				netopts = []string{"unixgram", "unix"}
			} else {
				netopts = []string{"tls"}
			}
		}
		addropts = []string{config.Address}
	} else {
		// This was copied from the standard log/syslog package, they do the same
		// and try to guess at runtime where and what kind of socket syslogd is
		// using.
		netopts = []string{"unixgram", "unix"}
		addropts = []string{"/dev/log", "/var/run/syslog", "/var/run/log"}
	}

connect:
	for _, n := range netopts {
		for _, a := range addropts {
			if backend, err = dialWriter(n, a, config.TLS); err == nil {
				break connect
			}
		}
	}

	if err != nil {
		return
	}

	w = newWriter(writerConfig{
		backend:    backend,
		template:   config.Template,
		timeFormat: config.TimeFormat,
	})
	return
}

type writerConfig struct {
	backend    io.Writer
	template   string
	timeFormat string
}

func newWriter(config writerConfig) *writer {
	var out func(*writer, message) error
	var flush func() error

	if len(config.timeFormat) == 0 {
		config.timeFormat = time.Stamp
	}

	if len(config.template) == 0 {
		config.template = DefaultTemplate
	}

	switch b := config.backend.(type) {
	case bufferedWriter:
		out, flush = (*writer).directWrite, b.Flush
	default:
		out, flush = (*writer).bufferedWrite, func() error { return nil }
	}

	return &writer{
		writerConfig: config,
		flush:        flush,
		out:          out,
		tpl:          newWriterTemplate(config.template),
	}
}

func newWriterTemplate(format string) *template.Template {
	if !strings.HasSuffix(format, "\n") {
		format += "\n"
	}
	t := template.New("syslog")
	template.Must(t.Parse(format))
	return t
}

type writer struct {
	writerConfig
	buf   bytes.Buffer
	tpl   *template.Template
	out   func(*writer, message) error
	flush func() error
}

func (w *writer) Close() (err error) {
	if c, ok := w.backend.(io.Closer); ok {
		err = c.Close()
	}
	return
}

func (w *writer) WriteMessageBatch(batch []ecslogs.Message) (err error) {
	for _, msg := range batch {
		if err = w.write(msg); err != nil {
			return
		}
	}
	return w.flush()
}

func (w *writer) WriteMessage(msg ecslogs.Message) (err error) {
	if err = w.write(msg); err == nil {
		err = w.flush()
	}
	return
}

func (w *writer) write(msg ecslogs.Message) (err error) {
	m := message{
		PRIVAL:    int(msg.Level) + 8, // +8 is for user-level messages facility
		HOSTNAME:  msg.Host,
		MSGID:     msg.ID,
		GROUP:     msg.Group,
		STREAM:    msg.Stream,
		MSG:       msg.Content,
		TIMESTAMP: msg.Time.Format(w.timeFormat),
	}

	if len(m.HOSTNAME) == 0 {
		m.HOSTNAME = "-"
	}

	if len(m.MSGID) == 0 {
		m.MSGID = "-"
	}

	if msg.PID == 0 {
		m.PROCID = "-"
	} else {
		m.PROCID = strconv.Itoa(msg.PID)
	}

	if len(msg.File) != 0 || len(msg.Func) != 0 {
		m.SOURCE = fmt.Sprintf("%s:%s:%d", msg.File, msg.Func, msg.Line)
	}

	return w.out(w, m)
}

func (w *writer) directWrite(m message) (err error) {
	return w.tpl.Execute(w.backend, m)
}

func (w *writer) bufferedWrite(m message) (err error) {
	w.buf.Reset()
	w.tpl.Execute(&w.buf, m)
	_, err = w.backend.Write(w.buf.Bytes())
	return
}

type message struct {
	PRIVAL    int
	HOSTNAME  string
	PROCID    string
	MSGID     string
	GROUP     string
	STREAM    string
	MSG       string
	SOURCE    string
	TIMESTAMP string
}

type bufferedWriter interface {
	Flush() error
}

type bufferedConn struct {
	buf  *bufio.Writer
	conn net.Conn
}

func (c bufferedConn) Close() error                { return c.conn.Close() }
func (c bufferedConn) Flush() error                { return c.buf.Flush() }
func (c bufferedConn) Write(b []byte) (int, error) { return c.buf.Write(b) }

func dialWriter(network string, address string, config *tls.Config) (w io.Writer, err error) {
	var conn net.Conn

	if network == "tls" {
		conn, err = tls.Dial("tcp", address, config)
	} else {
		conn, err = net.Dial(network, address)
	}

	if err == nil {
		switch network {
		case "udp", "udp4", "udp6", "unixgram", "unixpacket":
			w = conn
		default:
			w = bufferedConn{
				conn: conn,
				buf:  bufio.NewWriter(conn),
			}
		}
	}

	return
}
<commit_msg>better guess the syslog configuration<commit_after>package syslog

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/segmentio/ecs-logs/lib"
)

const (
	DefaultTemplate = "<{{.PRIVAL}}>{{.TIMESTAMP}} {{.GROUP}}[{{.STREAM}}]: {{.MSG}}"
)

type WriterConfig struct {
	Network    string
	Address    string
	Template   string
	TimeFormat string
	TLS        *tls.Config
}

func NewWriter(group string, stream string) (w ecslogs.Writer, err error) {
	var c WriterConfig
	var s string
	var u *url.URL

	if s = os.Getenv("SYSLOG_URL"); len(s) != 0 {
		if u, err = url.Parse(s); err != nil {
			err = fmt.Errorf("invalid syslog URL: %s", err)
			return
		}

		c.Network = u.Scheme
		c.Address = u.Host
	}

	c.Template = os.Getenv("SYSLOG_TEMPLATE")
	c.TimeFormat = os.Getenv("SYSLOG_TIME_FORMAT")

	return DialWriter(c)
}

func DialWriter(config WriterConfig) (w ecslogs.Writer, err error) {
	var netopts []string
	var addropts []string
	var backend io.Writer

	if len(config.Network) != 0 {
		netopts = []string{config.Network}
	} else {
		// When starting with a '/' we assume it's gonna be a file path,
		// otherwise we fallback to trying a TLS connection so we don't
		// implicitly send logs over an unsecured link.
		if len(config.Address) == 0 || strings.HasPrefix(config.Address, "/") {
			config.Network = "unix"
			netopts = []string{"unixgram", "unix"}
		} else {
			netopts = []string{"tls"}
		}
	}

	if len(config.Address) != 0 {
		addropts = []string{config.Address}
	} else if strings.HasPrefix(config.Network, "unix") {
		// This was copied from the standard log/syslog package, they do the same
		// and try to guess at runtime where and what kind of socket syslogd is
		// using.
		addropts = []string{"/dev/log", "/var/run/syslog", "/var/run/log"}
	} else {
		// The config doesn't point to a unix domain socket, falling back to trying
		// to connect to syslogd over a network interface.
		addropts = []string{"localhost:514"}
	}

connect:
	for _, n := range netopts {
		for _, a := range addropts {
			if backend, err = dialWriter(n, a, config.TLS); err == nil {
				break connect
			}
		}
	}

	if err != nil {
		return
	}

	w = newWriter(writerConfig{
		backend:    backend,
		template:   config.Template,
		timeFormat: config.TimeFormat,
	})
	return
}

type writerConfig struct {
	backend    io.Writer
	template   string
	timeFormat string
}

func newWriter(config writerConfig) *writer {
	var out func(*writer, message) error
	var flush func() error

	if len(config.timeFormat) == 0 {
		config.timeFormat = time.Stamp
	}

	if len(config.template) == 0 {
		config.template = DefaultTemplate
	}

	switch b := config.backend.(type) {
	case bufferedWriter:
		out, flush = (*writer).directWrite, b.Flush
	default:
		out, flush = (*writer).bufferedWrite, func() error { return nil }
	}

	return &writer{
		writerConfig: config,
		flush:        flush,
		out:          out,
		tpl:          newWriterTemplate(config.template),
	}
}

func newWriterTemplate(format string) *template.Template {
	if !strings.HasSuffix(format, "\n") {
		format += "\n"
	}
	t := template.New("syslog")
	template.Must(t.Parse(format))
	return t
}

type writer struct {
	writerConfig
	buf   bytes.Buffer
	tpl   *template.Template
	out   func(*writer, message) error
	flush func() error
}

func (w *writer) Close() (err error) {
	if c, ok := w.backend.(io.Closer); ok {
		err = c.Close()
	}
	return
}

func (w *writer) WriteMessageBatch(batch []ecslogs.Message) (err error) {
	for _, msg := range batch {
		if err = w.write(msg); err != nil {
			return
		}
	}
	return w.flush()
}

func (w *writer) WriteMessage(msg ecslogs.Message) (err error) {
	if err = w.write(msg); err == nil {
		err = w.flush()
	}
	return
}

func (w *writer) write(msg ecslogs.Message) (err error) {
	m := message{
		PRIVAL:    int(msg.Level) + 8, // +8 is for user-level messages facility
		HOSTNAME:  msg.Host,
		MSGID:     msg.ID,
		GROUP:     msg.Group,
		STREAM:    msg.Stream,
		MSG:       msg.Content,
		TIMESTAMP: msg.Time.Format(w.timeFormat),
	}

	if len(m.HOSTNAME) == 0 {
		m.HOSTNAME = "-"
	}

	if len(m.MSGID) == 0 {
		m.MSGID = "-"
	}

	if msg.PID == 0 {
		m.PROCID = "-"
	} else {
		m.PROCID = strconv.Itoa(msg.PID)
	}

	if len(msg.File) != 0 || len(msg.Func) != 0 {
		m.SOURCE = fmt.Sprintf("%s:%s:%d", msg.File, msg.Func, msg.Line)
	}

	return w.out(w, m)
}

func (w *writer) directWrite(m message) (err error) {
	return w.tpl.Execute(w.backend, m)
}

func (w *writer) bufferedWrite(m message) (err error) {
	w.buf.Reset()
	w.tpl.Execute(&w.buf, m)
	_, err = w.backend.Write(w.buf.Bytes())
	return
}

type message struct {
	PRIVAL    int
	HOSTNAME  string
	PROCID    string
	MSGID     string
	GROUP     string
	STREAM    string
	MSG       string
	SOURCE    string
	TIMESTAMP string
}

type bufferedWriter interface {
	Flush() error
}

type bufferedConn struct {
	buf  *bufio.Writer
	conn net.Conn
}

func (c bufferedConn) Close() error                { return c.conn.Close() }
func (c bufferedConn) Flush() error                { return c.buf.Flush() }
func (c bufferedConn) Write(b []byte) (int, error) { return c.buf.Write(b) }

func dialWriter(network string, address string, config *tls.Config) (w io.Writer, err error) {
	var conn net.Conn

	if network == "tls" {
		conn, err = tls.Dial("tcp", address, config)
	} else {
		conn, err = net.Dial(network, address)
	}

	if err == nil {
		switch network {
		case "udp", "udp4", "udp6", "unixgram", "unixpacket":
			w = conn
		default:
			w = bufferedConn{
				conn: conn,
				buf:  bufio.NewWriter(conn),
			}
		}
	}

	return
}
<|endoftext|>


<commit_before>package entry

import (
	"bufio"
	"github.com/ghthor/gospec"
	. "github.com/ghthor/gospec"
	"github.com/ghthor/journal/idea"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"time"
)

func init() {
	var err error
	if _, err = os.Stat("_test/"); os.IsNotExist(err) {
		err = os.Mkdir("_test/", 0755)
	}

	if err != nil {
		log.Fatal(err)
	}
}

func DescribeAnEntry(c gospec.Context) {
	td, err := ioutil.TempDir("_test", "entry_")
	c.Assume(err, IsNil)

	ne := New(td)

	c.Specify("an entry", func() {
		c.Specify("can be opened", func() {
			c.Specify("at a specific time", func() {
				t := time.Date(2006, time.January, 1, 1, 0, 0, 0, time.UTC)
				oe, err := ne.Open(func() time.Time {
					return t
				}, nil)
				c.Assume(err, IsNil)
				c.Expect(oe.OpenedAt(), Equals, t)
			})

			c.Specify("with a list of ideas", func() {
				ideas := []idea.Idea{{
					Name:   "Active Idea",
					Status: idea.IS_Active,
					Body:   "Some text\n",
				}, {
					Name:   "Another Idea",
					Status: idea.IS_Active,
					Body:   "Some other text\n",
				}}

				oe, err := ne.Open(time.Now, ideas)
				c.Assume(err, IsNil)
				for i, idea := range oe.Ideas() {
					c.Expect(idea, Equals, ideas[i])
				}
			})
		})
		t := time.Date(2006, time.January, 1, 1, 0, 0, 0, time.UTC)

		ideas := []idea.Idea{{
			Name:   "Active Idea",
			Status: idea.IS_Active,
			Body:   "Some text\n",
		}, {
			Name:   "Another Idea",
			Status: idea.IS_Active,
			Body:   "Some other text\n",
		}}

		c.Specify("that is open", func() {
			oe, err := ne.Open(func() time.Time {
				return t
			}, ideas)
			c.Assume(err, IsNil)
			c.Assume(oe.OpenedAt(), Equals, t)

			defer func() {
				_, _, err := oe.Close()
				c.Assume(err, IsNil)
				// Verify that the *os.File was closed
				c.Expect(oe.(*openEntry).file.Close(), Not(IsNil))
			}()

			filename := filepath.Join(td, t.Format(filenameLayout))

			c.Specify("is a file", func() {
				_, err := os.Stat(filename)
				c.Expect(os.IsNotExist(err), IsFalse)

				actualBytes, err := ioutil.ReadFile(filename)
				c.Expect(err, IsNil)
				c.Expect(string(actualBytes), Equals,
					`Sun Jan  1 01:00:00 UTC 2006

#~ Title(will be used as commit message)
TODO Make this some random quote or something stupid

## [active] Active Idea
Some text

## [active] Another Idea
Some other text
`)
			})

			c.Specify("will have the time opened as the first line of the entry", func() {
				f, err := os.OpenFile(filename, os.O_RDONLY, 0600)
				c.Assume(err, IsNil)
				defer f.Close()

				scanner := bufio.NewScanner(f)
				c.Assume(scanner.Scan(), IsTrue)
				c.Expect(scanner.Text(), Equals, t.Format(time.UnixDate))
			})

			c.Specify("will have a list of ideas appended to the entry", func() {
			})
			c.Specify("can be editted by a text editor", func() {
			})
			c.Specify("can be closed", func() {
			})
		})

		c.Specify("that is closed", func() {
			c.Specify("will have all ideas removed from the entry", func() {
			})
			c.Specify("will have the time closed as the last line of the entry", func() {
			})
			c.Specify("can be commited to the git repository", func() {
			})
		})
	})
}
<commit_msg>Specified that the Ideas are appended to the entries file<commit_after>package entry

import (
	"bufio"
	"github.com/ghthor/gospec"
	. "github.com/ghthor/gospec"
	"github.com/ghthor/journal/idea"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"time"
)

func init() {
	var err error
	if _, err = os.Stat("_test/"); os.IsNotExist(err) {
		err = os.Mkdir("_test/", 0755)
	}

	if err != nil {
		log.Fatal(err)
	}
}

func DescribeAnEntry(c gospec.Context) {
	td, err := ioutil.TempDir("_test", "entry_")
	c.Assume(err, IsNil)

	ne := New(td)

	c.Specify("an entry", func() {
		c.Specify("can be opened", func() {
			c.Specify("at a specific time", func() {
				t := time.Date(2006, time.January, 1, 1, 0, 0, 0, time.UTC)
				oe, err := ne.Open(func() time.Time {
					return t
				}, nil)
				c.Assume(err, IsNil)
				c.Expect(oe.OpenedAt(), Equals, t)
			})

			c.Specify("with a list of ideas", func() {
				ideas := []idea.Idea{{
					Name:   "Active Idea",
					Status: idea.IS_Active,
					Body:   "Some text\n",
				}, {
					Name:   "Another Idea",
					Status: idea.IS_Active,
					Body:   "Some other text\n",
				}}

				oe, err := ne.Open(time.Now, ideas)
				c.Assume(err, IsNil)
				for i, idea := range oe.Ideas() {
					c.Expect(idea, Equals, ideas[i])
				}
			})
		})
		t := time.Date(2006, time.January, 1, 1, 0, 0, 0, time.UTC)

		ideas := []idea.Idea{{
			Name:   "Active Idea",
			Status: idea.IS_Active,
			Body:   "Some text\n",
		}, {
			Name:   "Another Idea",
			Status: idea.IS_Active,
			Body:   "Some other text\n",
		}}

		c.Specify("that is open", func() {
			oe, err := ne.Open(func() time.Time {
				return t
			}, ideas)
			c.Assume(err, IsNil)
			c.Assume(oe.OpenedAt(), Equals, t)

			defer func() {
				_, _, err := oe.Close()
				c.Assume(err, IsNil)
				// Verify that the *os.File was closed
				c.Expect(oe.(*openEntry).file.Close(), Not(IsNil))
			}()

			filename := filepath.Join(td, t.Format(filenameLayout))

			c.Specify("is a file", func() {
				_, err := os.Stat(filename)
				c.Expect(os.IsNotExist(err), IsFalse)

				actualBytes, err := ioutil.ReadFile(filename)
				c.Expect(err, IsNil)
				c.Expect(string(actualBytes), Equals,
					`Sun Jan  1 01:00:00 UTC 2006

#~ Title(will be used as commit message)
TODO Make this some random quote or something stupid

## [active] Active Idea
Some text

## [active] Another Idea
Some other text
`)
			})

			f, err := os.OpenFile(filename, os.O_RDONLY, 0600)
			c.Assume(err, IsNil)
			defer f.Close()

			c.Specify("will have the time opened as the first line of the entry", func() {
				scanner := bufio.NewScanner(f)
				c.Assume(scanner.Scan(), IsTrue)
				c.Expect(scanner.Text(), Equals, t.Format(time.UnixDate))
			})

			c.Specify("will have a list of ideas appended to the entry", func() {
				scanner := idea.NewIdeaScanner(f)
				for i := 0; i < len(ideas); i++ {
					c.Assume(scanner.Scan(), IsTrue)
					c.Expect(*scanner.Idea(), Equals, ideas[i])
				}
			})

			c.Specify("can be editted by a text editor", func() {
			})
			c.Specify("can be closed", func() {
			})
		})

		c.Specify("that is closed", func() {
			c.Specify("will have all ideas removed from the entry", func() {
			})
			c.Specify("will have the time closed as the last line of the entry", func() {
			})
			c.Specify("can be commited to the git repository", func() {
			})
		})
	})
}
<|endoftext|>


<commit_before>package systemtokens

import (
	"errors"
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/rancher/rancher/pkg/auth/tokens"
	"github.com/rancher/rancher/pkg/randomtoken"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func NewSystemTokensFromScale(mgmt *config.ScaledContext) *SystemTokens {
	return &SystemTokens{
		tokenLister:  mgmt.Management.Tokens("").Controller().Lister(),
		tokenClient:  mgmt.Management.Tokens(""),
		tokenCache:   map[string]string{},
		isHA:         mgmt.PeerManager != nil,
		haIdentifier: getIdentifier(),
	}
}

type SystemTokens struct {
	isHA            bool
	tokenClient     v3.TokenInterface
	tokenLister     v3.TokenLister
	haIdentifier    string
	tokenCache      map[string]string
	tokenCacheMutex sync.RWMutex
}

// Multiple rancher management servers might attempt to use the same key name
// Because the token is hashed, any regeneration of an existing token name invalidates all previous tokens
// So to avoid those HA bugs, we need to uniquely create and re-use each token for each
// requesting management server, which we do here by generating a random string.
func getIdentifier() string {
	randString, err := randomtoken.Generate()
	if err != nil {
		return strconv.Itoa(rand.Intn(10000))
	}
	return randString[0:5]
}

// EnsureSystemToken gets or creates tokens for management use, and keeps them in memory inside contexts.
// Appends identifier to key name
// TTL defaults to 1 hour, after that this method will auto-refresh. If your token will be in use for more
// than one hour without calling this method again you must pass in an overrideTTL.
// However, the overrideTTL must not be 0, otherwise the token will never be cleaned up.
func (t *SystemTokens) EnsureSystemToken(key, description, kind, username string, overrideTTL *int64) (string, error) {
	if overrideTTL != nil && *overrideTTL == 0 {
		return "", errors.New("TTL for system token must not be zero") // no way to cleanup token
	}
	key = fmt.Sprintf("%s-%s", key, t.haIdentifier) // append hashed identifier, see getIdentifier
	token, err := t.tokenLister.Get("", key)
	if err != nil && !apierrors.IsNotFound(err) {
		return "", err
	}

	if err == nil && !tokens.IsExpired(*token) {
		t.tokenCacheMutex.RLock()
		val, ok := t.tokenCache[key]
		t.tokenCacheMutex.RUnlock()
		if ok {
			return val, nil
		}
	}
	// needs fresh token because its missing or expired
	val, err := t.createOrUpdateSystemToken(key, description, kind, username, overrideTTL)
	if err != nil {
		return "", err
	}

	t.tokenCacheMutex.Lock()
	defer t.tokenCacheMutex.Unlock()
	t.tokenCache[key] = val
	return val, nil
}

// Creates token obj with hashed token, returns token. Overwrites if pre-existing.
func (t *SystemTokens) createOrUpdateSystemToken(tokenName, description, kind, userName string, overrideTTL *int64) (string, error) {
	if strings.HasPrefix(tokenName, "token-") {
		return "", errors.New("token names can't start with token-")
	}
	key, err := randomtoken.Generate()
	if err != nil {
		return "", errors.New("failed to generate token key")
	}

	token, err := t.tokenLister.Get("", tokenName)
	if err != nil && !apierrors.IsNotFound(err) {
		return "", err
	}
	if token != nil {
		token.Token = key
		err = tokens.ConvertTokenKeyToHash(token)
		if err != nil {
			return "", err
		}
		logrus.Infof("Updating system token for %v, token: %v", userName, tokenName)
		token, err = t.tokenClient.Update(token)
		if err != nil {
			return "", err
		}
	} else {
		token = &v3.Token{
			ObjectMeta: v1.ObjectMeta{
				Name: tokenName,
				Labels: map[string]string{
					tokens.UserIDLabel:    userName,
					tokens.TokenKindLabel: kind,
				},
			},
			TTLMillis:    3600000, // 1 hour, token purge daemon will cleanup
			Description:  description,
			UserID:       userName,
			AuthProvider: "local",
			IsDerived:    true,
			Token:        key,
		}
		if overrideTTL != nil {
			token.TTLMillis = *overrideTTL
		}
		if t.isHA {
			if token.Annotations == nil {
				token.Annotations = make(map[string]string)
			}
			// For debugging purposes we set hostname as annotation, which in HA is the pod name
			token.Annotations[tokens.HAIdentifier] = os.Getenv("HOSTNAME")
		}

		err = tokens.ConvertTokenKeyToHash(token)
		if err != nil {
			return "", err
		}
		logrus.Infof("Creating system token for %v, token: %v", userName, tokenName)
		token, err = t.tokenClient.Create(token)
		if err != nil {
			return "", err
		}
	}
	return token.Name + ":" + key, nil
}
<commit_msg>Fixes system-token cache create bug<commit_after>package systemtokens

import (
	"errors"
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/rancher/rancher/pkg/auth/tokens"
	"github.com/rancher/rancher/pkg/randomtoken"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func NewSystemTokensFromScale(mgmt *config.ScaledContext) *SystemTokens {
	return &SystemTokens{
		tokenLister:  mgmt.Management.Tokens("").Controller().Lister(),
		tokenClient:  mgmt.Management.Tokens(""),
		tokenCache:   map[string]string{},
		isHA:         mgmt.PeerManager != nil,
		haIdentifier: getIdentifier(),
	}
}

type SystemTokens struct {
	isHA            bool
	tokenClient     v3.TokenInterface
	tokenLister     v3.TokenLister
	haIdentifier    string
	tokenCache      map[string]string
	tokenCacheMutex sync.RWMutex
}

// Multiple rancher management servers might attempt to use the same token name
// Because the token is hashed, any regeneration of an existing token name invalidates all previous tokens
// So to avoid those HA bugs, we need to uniquely create and re-use each token for each
// requesting management server, which we do here by generating a random string.
func getIdentifier() string {
	randString, err := randomtoken.Generate()
	if err != nil {
		return strconv.Itoa(rand.Intn(10000))
	}
	return randString[0:5]
}

// EnsureSystemToken gets or creates tokens for management use, and keeps them in memory inside contexts.
// Appends identifier to token name
// TTL defaults to 1 hour, after that this method will auto-refresh. If your token will be in use for more
// than one hour without calling this method again you must pass in an overrideTTL.
// However, the overrideTTL must not be 0, otherwise the token will never be cleaned up.
func (t *SystemTokens) EnsureSystemToken(tokenName, description, kind, username string, overrideTTL *int64) (string, error) {
	if overrideTTL != nil && *overrideTTL == 0 {
		return "", errors.New("TTL for system token must not be zero") // no way to cleanup token
	}
	tokenName = fmt.Sprintf("%s-%s", tokenName, t.haIdentifier) // append hashed identifier, see getIdentifier
	token, err := t.tokenLister.Get("", tokenName)
	if err != nil && !apierrors.IsNotFound(err) {
		return "", err
	}

	if err == nil && !tokens.IsExpired(*token) {
		if tokenVal := t.getTokenValue(tokenName); tokenVal != "" {
			return tokenVal, nil
		}
	}
	// needs fresh token because its missing or expired
	val, err := t.createOrUpdateSystemToken(tokenName, description, kind, username, overrideTTL)
	if err != nil {
		return "", err
	}
	return val, nil
}

func (t *SystemTokens) getTokenValue(tokenName string) string {
	t.tokenCacheMutex.RLock()
	val, ok := t.tokenCache[tokenName]
	t.tokenCacheMutex.RUnlock()
	if ok {
		return val
	}
	return ""
}

// Creates token obj with hashed token, returns token. Overwrites if pre-existing.
func (t *SystemTokens) createOrUpdateSystemToken(tokenName, description, kind, userName string, overrideTTL *int64) (string, error) {
	if strings.HasPrefix(tokenName, "token-") {
		return "", errors.New("token names can't start with token-")
	}
	key, err := randomtoken.Generate()
	if err != nil {
		return "", errors.New("failed to generate token key")
	}

	token, err := t.tokenLister.Get("", tokenName)
	if err != nil {
		if !apierrors.IsNotFound(err) {
			return "", err
		}
		// not found error, make new token
		token = &v3.Token{
			ObjectMeta: v1.ObjectMeta{
				Name: tokenName,
				Labels: map[string]string{
					tokens.UserIDLabel:    userName,
					tokens.TokenKindLabel: kind,
				},
			},
			TTLMillis:    3600000, // 1 hour, token purge daemon will cleanup
			Description:  description,
			UserID:       userName,
			AuthProvider: "local",
			IsDerived:    true,
			Token:        key,
		}
		if overrideTTL != nil {
			token.TTLMillis = *overrideTTL
		}
		if t.isHA {
			if token.Annotations == nil {
				token.Annotations = make(map[string]string)
			}
			// For debugging purposes we set hostname as annotation, which in HA is the pod name
			token.Annotations[tokens.HAIdentifier] = os.Getenv("HOSTNAME")
		}

		err = tokens.ConvertTokenKeyToHash(token)
		if err != nil {
			return "", err
		}
		logrus.Infof("Creating system token for %v, token: %v", userName, tokenName)
		token, err = t.tokenClient.Create(token)
		if err != nil {
			// race condition on create, someone else just created so return that val instead
			if apierrors.IsAlreadyExists(err) {
				if tokenVal := t.getTokenValue(tokenName); tokenVal != "" {
					return tokenVal, nil
				}
				return "", fmt.Errorf("system token not found in cache: %s", tokenName)
			}
			return "", err
		}
	} else {
		token.Token = key
		err = tokens.ConvertTokenKeyToHash(token)
		if err != nil {
			return "", err
		}
		logrus.Infof("Updating system token for %v, token: %v", userName, tokenName)
		token, err = t.tokenClient.Update(token)
		if err != nil {
			return "", err
		}
	}
	// Fresh token needs saving in cache
	fullVal := fmt.Sprintf("%s:%s", token.Name, key)
	t.tokenCacheMutex.Lock()
	defer t.tokenCacheMutex.Unlock()
	t.tokenCache[tokenName] = fullVal
	return fullVal, nil
}
<|endoftext|>