<commit_before>// Copyright 2017 The OpenSDS Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/*
This module implements the entry into operations of storageDock module.

*/

package discovery

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"time"

	log "github.com/golang/glog"
	"github.com/opensds/opensds/contrib/connector"
	"github.com/opensds/opensds/contrib/drivers"
	fd "github.com/opensds/opensds/contrib/drivers/filesharedrivers"
	"github.com/opensds/opensds/contrib/drivers/utils/config"
	c "github.com/opensds/opensds/pkg/context"
	"github.com/opensds/opensds/pkg/db"
	"github.com/opensds/opensds/pkg/model"
	"github.com/opensds/opensds/pkg/utils"
	. "github.com/opensds/opensds/pkg/utils/config"
	uuid "github.com/satori/go.uuid"
)

const (
	availableStatus   = "available"
	unavailableStatus = "unavailable"
)

type Context struct {
	StopChan chan bool
	ErrChan  chan error
	MetaChan chan string
}

func DiscoveryAndReport(dd DockDiscoverer, ctx *Context) {
	for {
		select {
		case <-ctx.StopChan:
			return
		default:
			if err := dd.Discover(); err != nil {
				ctx.ErrChan <- err
			}

			if err := dd.Report(); err != nil {
				ctx.ErrChan <- err
			}
		}

		time.Sleep(60 * time.Second)
	}
}

type DockDiscoverer interface {
	Init() error

	Discover() error

	Report() error
}

// NewDockDiscoverer method creates a new DockDiscoverer.
func NewDockDiscoverer(dockType string) DockDiscoverer {
	switch dockType {
	case model.DockTypeProvioner:
		return &provisionDockDiscoverer{
			DockRegister: NewDockRegister(),
		}
	case model.DockTypeAttacher:
		return &attachDockDiscoverer{
			DockRegister: NewDockRegister(),
		}
	}
	return nil
}

// provisionDockDiscoverer is a struct for exposing some operations of provision
// dock service discovery.
type provisionDockDiscoverer struct {
	*DockRegister

	dcks []*model.DockSpec
	pols []*model.StoragePoolSpec
}

func (pdd *provisionDockDiscoverer) Init() error {
	// Load resource from specified file
	bm := GetBackendsMap()
	host, err := os.Hostname()
	if err != nil {
		log.Error("When get os hostname:", err)
		return err
	}

	for _, v := range CONF.EnabledBackends {
		b := bm[v]
		if b.Name == "" {
			continue
		}

		dck := &model.DockSpec{
			BaseModel: &model.BaseModel{
				Id: uuid.NewV5(uuid.NamespaceOID, host+":"+b.DriverName).String(),
			},
			Name:        b.Name,
			Description: b.Description,
			DriverName:  b.DriverName,
			Endpoint:    CONF.OsdsDock.ApiEndpoint,
			NodeId:      host,
			Type:        model.DockTypeProvioner,
			Metadata:    map[string]string{"HostReplicationDriver": CONF.OsdsDock.HostBasedReplicationDriver},
		}
		// Update the id if the dock is already in etcd
		name := map[string][]string{
			"Name": {dck.Name},
		}
		docks, err := pdd.DockRegister.c.ListDocksWithFilter(c.NewAdminContext(), name)
		if err == nil && len(docks) != 0 {
			dck.Id = docks[0].Id
		}
		pdd.dcks = append(pdd.dcks, dck)
	}

	return nil
}

var filesharedrivers = []string{config.NFSDriverType, config.HuaweiOceanStorFileDriverType, config.ManilaDriverType, config.ChubaofsDriverType}

func (pdd *provisionDockDiscoverer) Discover() error {
	// Clear existing pool info
	pdd.pols = pdd.pols[:0]
	var pols []*model.StoragePoolSpec
	var err error
	var polsInDb []*model.StoragePoolSpec
	ctx := c.NewAdminContext()
	polsInDb, err = pdd.c.ListPools(ctx)
	if err != nil {
		return fmt.Errorf("can not read pools in db")
	}
	dbPolsMap := make(map[string]map[string]*model.StoragePoolSpec)
	for _, dck := range pdd.dcks {
		dbPolsMap[dck.Id] = make(map[string]*model.StoragePoolSpec)
	}
	for _, polInDb := range polsInDb {
		if dbPolsMap[polInDb.DockId] != nil {
			polInDb.Status = unavailableStatus
			dbPolsMap[polInDb.DockId][polInDb.Id] = polInDb
		}
	}
	for _, dck := range pdd.dcks {
		// Call function of StorageDrivers configured by storage drivers.
		if utils.Contains(filesharedrivers, dck.DriverName) {
			d := fd.Init(dck.DriverName)
			defer fd.Clean(d)
			pols, err = d.ListPools()
			for _, pol := range pols {
				log.Infof("Backend %s discovered pool %s", dck.DriverName, pol.Name)
				delete(dbPolsMap[dck.Id], pol.Id)
				pol.DockId = dck.Id
				pol.Status = availableStatus
			}
		} else {
			d := drivers.Init(dck.DriverName)
			defer drivers.Clean(d)
			pols, err = d.ListPools()

			replicationDriverName := dck.Metadata["HostReplicationDriver"]
			replicationType := model.ReplicationTypeHost
			if drivers.IsSupportArrayBasedReplication(dck.DriverName) {
				replicationType = model.ReplicationTypeArray
				replicationDriverName = dck.DriverName
			}
			for _, pol := range pols {
				log.Infof("Backend %s discovered pool %s", dck.DriverName, pol.Name)
				name := map[string][]string{
					"Name":   {pol.Name},
					"dockId": {pol.DockId},
				}
				pools, err := pdd.c.ListPoolsWithFilter(ctx, name)
				if err == nil && len(pools) != 0 {
					pol.Id = pools[0].Id
				}

				delete(dbPolsMap[dck.Id], pol.Id)
				pol.DockId = dck.Id
				pol.ReplicationType = replicationType
				pol.ReplicationDriverName = replicationDriverName
				pol.Status = availableStatus
			}
		}
		if err != nil {
			log.Error("Call driver to list pools failed:", err)
			continue
		}

		if len(pols) == 0 {
			log.Warningf("The pool of dock %s is empty!\n", dck.Id)
		}

		pdd.pols = append(pdd.pols, pols...)
		for _, pol := range dbPolsMap[dck.Id] {
			pdd.pols = append(pdd.pols, pol)
		}

	}
	if len(pdd.pols) == 0 {
		return fmt.Errorf("there is no pool can be found")
	}

	return nil
}

func (pdd *provisionDockDiscoverer) Report() error {
	var err error

	// Store dock resources in database.
	for _, dck := range pdd.dcks {
		if err = pdd.Register(dck); err != nil {
			break
		}
	}

	// Store pool resources in database.
	for _, pol := range pdd.pols {
		if err != nil {
			break
		}
		err = pdd.Register(pol)
	}

	return err
}

// attachDockDiscoverer is a struct for exposing some operations of attach
// dock service discovery.
type attachDockDiscoverer struct {
	*DockRegister

	dck *model.DockSpec
}

func (add *attachDockDiscoverer) Init() error { return nil }

func (add *attachDockDiscoverer) Discover() error {
	host, err := os.Hostname()
	if err != nil {
		log.Error("When get os hostname:", err)
		return err
	}

	localIqn, err := connector.NewConnector(connector.IscsiDriver).GetInitiatorInfo()
	if err != nil {
		log.Warning("get initiator failed, ", err)
	}

	bindIp := CONF.BindIp
	if bindIp == "" {
		bindIp = connector.GetHostIP()
	}

	fcInitiator, err := connector.NewConnector(connector.FcDriver).GetInitiatorInfo()
	if err != nil {
		log.Warning("get initiator failed, ", err)
	}

	var wwpns []string
	for _, v := range strings.Split(fcInitiator, ",") {
		if strings.Contains(v, "node_name") {
			wwpns = append(wwpns, strings.Split(v, ":")[1])
		}
	}

	segments := strings.Split(CONF.OsdsDock.ApiEndpoint, ":")
	endpointIp := segments[len(segments)-2]
	add.dck = &model.DockSpec{
		BaseModel: &model.BaseModel{
			Id: uuid.NewV5(uuid.NamespaceOID, host+":"+endpointIp).String(),
		},
		Endpoint: CONF.OsdsDock.ApiEndpoint,
		NodeId:   host,
		Type:     model.DockTypeAttacher,
		Metadata: map[string]string{
			"Platform":  runtime.GOARCH,
			"OsType":    runtime.GOOS,
			"HostIp":    bindIp,
			"Initiator": localIqn,
			"WWPNS":     strings.Join(wwpns, ","),
		},
	}
	return nil
}

func (add *attachDockDiscoverer) Report() error {
	return add.Register(add.dck)
}

func NewDockRegister() *DockRegister {
	return &DockRegister{c: db.C}
}

type DockRegister struct {
	c db.Client
}

func (dr *DockRegister) Register(in interface{}) error {
	ctx := c.NewAdminContext()

	switch in.(type) {
	case *model.DockSpec:
		dck := in.(*model.DockSpec)
		// Call db module to create dock resource.
		if _, err := dr.c.CreateDock(ctx, dck); err != nil {
			log.Errorf("When create dock %s in db: %v\n", dck.Id, err)
			return err
		}
		break
	case *model.StoragePoolSpec:
		pol := in.(*model.StoragePoolSpec)
		// Call db module to create pool resource.
		if _, err := dr.c.CreatePool(ctx, pol); err != nil {
			log.Errorf("When create pool %s in db: %v\n", pol.Id, err)
			return err
		}
		break
	default:
		return fmt.Errorf("Resource type is not supported!")
	}

	return nil
}

func (dr *DockRegister) Unregister(in interface{}) error {
	ctx := c.NewAdminContext()

	switch in.(type) {
	case *model.DockSpec:
		dck := in.(*model.DockSpec)
		// Call db module to delete dock resource.
		if err := dr.c.DeleteDock(ctx, dck.Id); err != nil {
			log.Errorf("When delete dock %s in db: %v\n", dck.Id, err)
			return err
		}
		break
	case *model.StoragePoolSpec:
		pol := in.(*model.StoragePoolSpec)
		// Call db module to delete pool resource.
		if err := dr.c.DeletePool(ctx, pol.Id); err != nil {
			log.Errorf("When delete pool %s in db: %v\n", pol.Id, err)
			return err
		}
		break
	default:
		return fmt.Errorf("Resource type is not supported!")
	}

	return nil
}
<commit_msg>Use another dock id<commit_after>// Copyright 2017 The OpenSDS Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/*
This module implements the entry into operations of storageDock module.

*/

package discovery

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"time"

	log "github.com/golang/glog"
	"github.com/opensds/opensds/contrib/connector"
	"github.com/opensds/opensds/contrib/drivers"
	fd "github.com/opensds/opensds/contrib/drivers/filesharedrivers"
	"github.com/opensds/opensds/contrib/drivers/utils/config"
	c "github.com/opensds/opensds/pkg/context"
	"github.com/opensds/opensds/pkg/db"
	"github.com/opensds/opensds/pkg/model"
	"github.com/opensds/opensds/pkg/utils"
	. "github.com/opensds/opensds/pkg/utils/config"
	uuid "github.com/satori/go.uuid"
)

const (
	availableStatus   = "available"
	unavailableStatus = "unavailable"
)

type Context struct {
	StopChan chan bool
	ErrChan  chan error
	MetaChan chan string
}

func DiscoveryAndReport(dd DockDiscoverer, ctx *Context) {
	for {
		select {
		case <-ctx.StopChan:
			return
		default:
			if err := dd.Discover(); err != nil {
				ctx.ErrChan <- err
			}

			if err := dd.Report(); err != nil {
				ctx.ErrChan <- err
			}
		}

		time.Sleep(60 * time.Second)
	}
}

type DockDiscoverer interface {
	Init() error

	Discover() error

	Report() error
}

// NewDockDiscoverer method creates a new DockDiscoverer.
func NewDockDiscoverer(dockType string) DockDiscoverer {
	switch dockType {
	case model.DockTypeProvioner:
		return &provisionDockDiscoverer{
			DockRegister: NewDockRegister(),
		}
	case model.DockTypeAttacher:
		return &attachDockDiscoverer{
			DockRegister: NewDockRegister(),
		}
	}
	return nil
}

// provisionDockDiscoverer is a struct for exposing some operations of provision
// dock service discovery.
type provisionDockDiscoverer struct {
	*DockRegister

	dcks []*model.DockSpec
	pols []*model.StoragePoolSpec
}

func (pdd *provisionDockDiscoverer) Init() error {
	// Load resource from specified file
	bm := GetBackendsMap()
	host, err := os.Hostname()
	if err != nil {
		log.Error("When get os hostname:", err)
		return err
	}

	for _, v := range CONF.EnabledBackends {
		b := bm[v]
		if b.Name == "" {
			continue
		}

		dck := &model.DockSpec{
			BaseModel: &model.BaseModel{
				Id: uuid.NewV5(uuid.NamespaceOID, host+":"+b.DriverName).String(),
			},
			Name:        b.Name,
			Description: b.Description,
			DriverName:  b.DriverName,
			Endpoint:    CONF.OsdsDock.ApiEndpoint,
			NodeId:      host,
			Type:        model.DockTypeProvioner,
			Metadata:    map[string]string{"HostReplicationDriver": CONF.OsdsDock.HostBasedReplicationDriver},
		}
		// Update the id if the dock is already in etcd
		name := map[string][]string{
			"Name": {dck.Name},
		}
		docks, err := pdd.DockRegister.c.ListDocksWithFilter(c.NewAdminContext(), name)
		if err == nil && len(docks) != 0 {
			dck.Id = docks[0].Id
		}
		pdd.dcks = append(pdd.dcks, dck)
	}

	return nil
}

var filesharedrivers = []string{config.NFSDriverType, config.HuaweiOceanStorFileDriverType, config.ManilaDriverType, config.ChubaofsDriverType}

func (pdd *provisionDockDiscoverer) Discover() error {
	// Clear existing pool info
	pdd.pols = pdd.pols[:0]
	var pols []*model.StoragePoolSpec
	var err error
	var polsInDb []*model.StoragePoolSpec
	ctx := c.NewAdminContext()
	polsInDb, err = pdd.c.ListPools(ctx)
	if err != nil {
		return fmt.Errorf("can not read pools in db")
	}
	dbPolsMap := make(map[string]map[string]*model.StoragePoolSpec)
	for _, dck := range pdd.dcks {
		dbPolsMap[dck.Id] = make(map[string]*model.StoragePoolSpec)
	}
	for _, polInDb := range polsInDb {
		if dbPolsMap[polInDb.DockId] != nil {
			polInDb.Status = unavailableStatus
			dbPolsMap[polInDb.DockId][polInDb.Id] = polInDb
		}
	}
	for _, dck := range pdd.dcks {
		// Call function of StorageDrivers configured by storage drivers.
		if utils.Contains(filesharedrivers, dck.DriverName) {
			d := fd.Init(dck.DriverName)
			defer fd.Clean(d)
			pols, err = d.ListPools()
			for _, pol := range pols {
				log.Infof("Backend %s discovered pool %s", dck.DriverName, pol.Name)
				delete(dbPolsMap[dck.Id], pol.Id)
				pol.DockId = dck.Id
				pol.Status = availableStatus
			}
		} else {
			d := drivers.Init(dck.DriverName)
			defer drivers.Clean(d)
			pols, err = d.ListPools()

			replicationDriverName := dck.Metadata["HostReplicationDriver"]
			replicationType := model.ReplicationTypeHost
			if drivers.IsSupportArrayBasedReplication(dck.DriverName) {
				replicationType = model.ReplicationTypeArray
				replicationDriverName = dck.DriverName
			}
			for _, pol := range pols {
				log.Infof("Backend %s discovered pool %s", dck.DriverName, pol.Name)
				name := map[string][]string{
					"Name":   {pol.Name},
					"DockId": {dck.Id},
				}
				pools, err := pdd.c.ListPoolsWithFilter(ctx, name)
				if err == nil && len(pools) != 0 {
					pol.Id = pools[0].Id
				}

				delete(dbPolsMap[dck.Id], pol.Id)
				pol.DockId = dck.Id
				pol.ReplicationType = replicationType
				pol.ReplicationDriverName = replicationDriverName
				pol.Status = availableStatus
			}
		}
		if err != nil {
			log.Error("Call driver to list pools failed:", err)
			continue
		}

		if len(pols) == 0 {
			log.Warningf("The pool of dock %s is empty!\n", dck.Id)
		}

		pdd.pols = append(pdd.pols, pols...)
		for _, pol := range dbPolsMap[dck.Id] {
			pdd.pols = append(pdd.pols, pol)
		}

	}
	if len(pdd.pols) == 0 {
		return fmt.Errorf("there is no pool can be found")
	}

	return nil
}

func (pdd *provisionDockDiscoverer) Report() error {
	var err error

	// Store dock resources in database.
	for _, dck := range pdd.dcks {
		if err = pdd.Register(dck); err != nil {
			break
		}
	}

	// Store pool resources in database.
	for _, pol := range pdd.pols {
		if err != nil {
			break
		}
		err = pdd.Register(pol)
	}

	return err
}

// attachDockDiscoverer is a struct for exposing some operations of attach
// dock service discovery.
type attachDockDiscoverer struct {
	*DockRegister

	dck *model.DockSpec
}

func (add *attachDockDiscoverer) Init() error { return nil }

func (add *attachDockDiscoverer) Discover() error {
	host, err := os.Hostname()
	if err != nil {
		log.Error("When get os hostname:", err)
		return err
	}

	localIqn, err := connector.NewConnector(connector.IscsiDriver).GetInitiatorInfo()
	if err != nil {
		log.Warning("get initiator failed, ", err)
	}

	bindIp := CONF.BindIp
	if bindIp == "" {
		bindIp = connector.GetHostIP()
	}

	fcInitiator, err := connector.NewConnector(connector.FcDriver).GetInitiatorInfo()
	if err != nil {
		log.Warning("get initiator failed, ", err)
	}

	var wwpns []string
	for _, v := range strings.Split(fcInitiator, ",") {
		if strings.Contains(v, "node_name") {
			wwpns = append(wwpns, strings.Split(v, ":")[1])
		}
	}

	segments := strings.Split(CONF.OsdsDock.ApiEndpoint, ":")
	endpointIp := segments[len(segments)-2]
	add.dck = &model.DockSpec{
		BaseModel: &model.BaseModel{
			Id: uuid.NewV5(uuid.NamespaceOID, host+":"+endpointIp).String(),
		},
		Endpoint: CONF.OsdsDock.ApiEndpoint,
		NodeId:   host,
		Type:     model.DockTypeAttacher,
		Metadata: map[string]string{
			"Platform":  runtime.GOARCH,
			"OsType":    runtime.GOOS,
			"HostIp":    bindIp,
			"Initiator": localIqn,
			"WWPNS":     strings.Join(wwpns, ","),
		},
	}
	return nil
}

func (add *attachDockDiscoverer) Report() error {
	return add.Register(add.dck)
}

func NewDockRegister() *DockRegister {
	return &DockRegister{c: db.C}
}

type DockRegister struct {
	c db.Client
}

func (dr *DockRegister) Register(in interface{}) error {
	ctx := c.NewAdminContext()

	switch in.(type) {
	case *model.DockSpec:
		dck := in.(*model.DockSpec)
		// Call db module to create dock resource.
		if _, err := dr.c.CreateDock(ctx, dck); err != nil {
			log.Errorf("When create dock %s in db: %v\n", dck.Id, err)
			return err
		}
		break
	case *model.StoragePoolSpec:
		pol := in.(*model.StoragePoolSpec)
		// Call db module to create pool resource.
		if _, err := dr.c.CreatePool(ctx, pol); err != nil {
			log.Errorf("When create pool %s in db: %v\n", pol.Id, err)
			return err
		}
		break
	default:
		return fmt.Errorf("Resource type is not supported!")
	}

	return nil
}

func (dr *DockRegister) Unregister(in interface{}) error {
	ctx := c.NewAdminContext()

	switch in.(type) {
	case *model.DockSpec:
		dck := in.(*model.DockSpec)
		// Call db module to delete dock resource.
		if err := dr.c.DeleteDock(ctx, dck.Id); err != nil {
			log.Errorf("When delete dock %s in db: %v\n", dck.Id, err)
			return err
		}
		break
	case *model.StoragePoolSpec:
		pol := in.(*model.StoragePoolSpec)
		// Call db module to delete pool resource.
		if err := dr.c.DeletePool(ctx, pol.Id); err != nil {
			log.Errorf("When delete pool %s in db: %v\n", pol.Id, err)
			return err
		}
		break
	default:
		return fmt.Errorf("Resource type is not supported!")
	}

	return nil
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package endpoint

import (
	"context"
	"sort"
	"strings"
	"time"

	"github.com/cilium/cilium/api/v1/models"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/identity/cache"
	identitymodel "github.com/cilium/cilium/pkg/identity/model"
	cilium_v2 "github.com/cilium/cilium/pkg/k8s/apis/cilium.io/v2"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/node"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/policy"
)

// StatusNA is value of fields in the output of 'kubectl get cep' in case of disabled "--endpoint-status"
const (
	EndpointPolicyStateEnforcing    cilium_v2.EndpointPolicyState = "enforcing"
	EndpointPolicyStateNonEnforcing cilium_v2.EndpointPolicyState = "non-enforcing"
	EndpointPolicyStateDisabled     cilium_v2.EndpointPolicyState = "disabled"
)

func getEndpointStatusControllers(mdlControllers models.ControllerStatuses) (controllers cilium_v2.ControllerList) {
	for _, c := range mdlControllers {
		if c.Status == nil {
			continue
		}

		if c.Status.ConsecutiveFailureCount > 0 {
			s := cilium_v2.ControllerStatus{
				Configuration: c.Configuration,
				Name:          c.Name,
				UUID:          string(c.UUID),
				Status: cilium_v2.ControllerStatusStatus{
					ConsecutiveFailureCount: c.Status.ConsecutiveFailureCount,
					FailureCount:            c.Status.FailureCount,
					LastFailureMsg:          c.Status.LastFailureMsg,
					LastFailureTimestamp:    c.Status.LastFailureTimestamp.String(),
					LastSuccessTimestamp:    c.Status.LastSuccessTimestamp.String(),
					SuccessCount:            c.Status.SuccessCount,
				},
			}
			if controllers == nil {
				controllers = cilium_v2.ControllerList{s}
			} else {
				controllers = append(controllers, s)
			}
		}
	}

	if controllers != nil {
		controllers.Sort()
	}

	return
}

func (e *Endpoint) getEndpointStatusLog() (log []*models.EndpointStatusChange) {
	added := 0

	if s := e.status; s != nil {
		s.indexMU.RLock()
		defer s.indexMU.RUnlock()

		for i := s.lastIndex(); ; i-- {
			if i < 0 {
				i = maxLogs - 1
			}
			if i < len(s.Log) && s.Log[i] != nil {
				l := &models.EndpointStatusChange{
					Timestamp: s.Log[i].Timestamp.Format(time.RFC3339),
					Code:      s.Log[i].Status.Code.String(),
					Message:   s.Log[i].Status.Msg,
					State:     models.EndpointState(s.Log[i].Status.State),
				}

				if strings.ToLower(l.Code) != models.EndpointStatusChangeCodeOk {
					if log == nil {
						log = []*models.EndpointStatusChange{l}
					} else {
						log = append(log, l)
					}

					// Limit the number of endpoint log
					// entries to keep the size of the
					// EndpointStatus low.
					added++
					if added >= cilium_v2.EndpointStatusLogEntries {
						break
					}
				}
			}
			if i == s.Index {
				break
			}
		}
	}
	return
}

func getEndpointIdentity(mdlIdentity *models.Identity) (identity *cilium_v2.EndpointIdentity) {
	if mdlIdentity == nil {
		return
	}
	identity = &cilium_v2.EndpointIdentity{
		ID: mdlIdentity.ID,
	}

	identity.Labels = make([]string, len(mdlIdentity.Labels))
	copy(identity.Labels, mdlIdentity.Labels)
	sort.Strings(identity.Labels)
	return
}

func getEndpointNetworking(mdlNetworking *models.EndpointNetworking) (networking *cilium_v2.EndpointNetworking) {
	if mdlNetworking == nil {
		return nil
	}
	networking = &cilium_v2.EndpointNetworking{
		Addressing: make(cilium_v2.AddressPairList, len(mdlNetworking.Addressing)),
	}

	if option.Config.EnableIPv4 {
		networking.NodeIP = node.GetIPv4().String()
	} else {
		networking.NodeIP = node.GetIPv6().String()
	}

	for i, pair := range mdlNetworking.Addressing {
		networking.Addressing[i] = &cilium_v2.AddressPair{
			IPV4: pair.IPV4,
			IPV6: pair.IPV6,
		}
	}

	networking.Addressing.Sort()
	return
}

// updateLabels inserts the labels correnspoding to the specified identity into
// the AllowedIdentityTuple.
func updateLabels(allocator cache.IdentityAllocator, allowedIdentityTuple *cilium_v2.IdentityTuple, secID identity.NumericIdentity) {
	// IdentityUnknown denotes that this is an L4-only BPF
	// allow, so it applies to all identities. In this case
	// we should skip resolving the labels, because the
	// value 0 does not denote an allow for the "unknown"
	// identity, but instead an allow of all identities for
	// that port.
	if secID != identity.IdentityUnknown {
		identity := allocator.LookupIdentityByID(context.TODO(), secID)
		if identity != nil {
			var l labels.Labels
			if identity.CIDRLabel != nil {
				l = identity.CIDRLabel
			} else {
				l = identity.Labels
			}

			allowedIdentityTuple.IdentityLabels = l.StringMap()
		}
	}
}

// populateResponseWithPolicyKey inserts an AllowedIdentityTuple element into
// 'policy' which corresponds to the specified 'desiredPolicy'. If 'isDeny' is
// true, it will insert the policyKey into the 'Denied'.
func populateResponseWithPolicyKey(
	allocator cache.IdentityAllocator,
	policy *cilium_v2.EndpointPolicy,
	policyKey *policy.Key,
	isDeny bool,
) {
	identityTuple := cilium_v2.IdentityTuple{
		DestPort: policyKey.DestPort,
		Protocol: policyKey.Nexthdr,
		Identity: uint64(policyKey.Identity),
	}

	secID := identity.NumericIdentity(policyKey.Identity)
	updateLabels(allocator, &identityTuple, secID)

	switch {
	case policyKey.IsIngress():
		if isDeny {
			if policy.Ingress.Denied == nil {
				policy.Ingress.Denied = cilium_v2.DenyIdentityList{identityTuple}
			} else {
				policy.Ingress.Denied = append(policy.Ingress.Denied, identityTuple)
			}
		} else {
			if policy.Ingress.Allowed == nil {
				policy.Ingress.Allowed = cilium_v2.AllowedIdentityList{identityTuple}
			} else {
				policy.Ingress.Allowed = append(policy.Ingress.Allowed, identityTuple)
			}
		}
	case policyKey.IsEgress():
		if isDeny {
			if policy.Egress.Denied == nil {
				policy.Egress.Denied = cilium_v2.DenyIdentityList{identityTuple}
			} else {
				policy.Egress.Denied = append(policy.Egress.Denied, identityTuple)
			}
		} else {
			if policy.Egress.Allowed == nil {
				policy.Egress.Allowed = cilium_v2.AllowedIdentityList{identityTuple}
			} else {
				policy.Egress.Allowed = append(policy.Egress.Allowed, identityTuple)
			}
		}
	}
}

// getEndpointPolicy returns an API representation of the policy that the
// received Endpoint intends to apply.
func (e *Endpoint) getEndpointPolicy() (ep *cilium_v2.EndpointPolicy) {
	if e.desiredPolicy == nil {
		return
	}
	ep = &cilium_v2.EndpointPolicy{
		Ingress: &cilium_v2.EndpointPolicyDirection{
			Enforcing: !e.Options.IsEnabled(option.PolicyAuditMode) &&
				e.desiredPolicy.IngressPolicyEnabled,
		},
		Egress: &cilium_v2.EndpointPolicyDirection{
			Enforcing: !e.Options.IsEnabled(option.PolicyAuditMode) &&
				e.desiredPolicy.EgressPolicyEnabled,
		},
	}
	if ep.Ingress.Enforcing {
		ep.Ingress.State = EndpointPolicyStateEnforcing
	} else {
		ep.Ingress.State = EndpointPolicyStateNonEnforcing
	}
	if ep.Egress.Enforcing {
		ep.Egress.State = EndpointPolicyStateEnforcing
	} else {
		ep.Egress.State = EndpointPolicyStateNonEnforcing
	}

	// Handle allow-all cases
	allowsAllIngress, allowsAllEgress := e.desiredPolicy.AllowsIdentity(identity.IdentityUnknown)
	if allowsAllIngress {
		ep.Ingress.Allowed = cilium_v2.AllowedIdentityList{{}}
		ep.Ingress.Denied = cilium_v2.DenyIdentityList{{}}
	}
	if allowsAllEgress {
		ep.Egress.Allowed = cilium_v2.AllowedIdentityList{{}}
		ep.Egress.Denied = cilium_v2.DenyIdentityList{{}}
	}

	// If either ingress or egress policy is enabled, go through
	// the desired policy to populate the values.
	if !allowsAllIngress || !allowsAllEgress {
		allowsWorldIngress, allowsWorldEgress := e.desiredPolicy.AllowsIdentity(identity.ReservedIdentityWorld)

		for policyKey, policyValue := range e.desiredPolicy.PolicyMapState {
			// Skip listing identities if enforcement is disabled in direction,
			// or if the identity corresponds to a CIDR identity and the world is allowed.
			id := identity.NumericIdentity(policyKey.Identity)
			switch {
			case policyKey.IsIngress():
				if allowsAllIngress || (id.HasLocalScope() && allowsWorldIngress) {
					continue
				}
			case policyKey.IsEgress():
				if allowsAllEgress || (id.HasLocalScope() && allowsWorldEgress) {
					continue
				}
			}

			populateResponseWithPolicyKey(e.allocator, ep, &policyKey, policyValue.IsDeny)
		}
	}

	if ep.Ingress.Allowed != nil {
		ep.Ingress.Allowed.Sort()
	}
	if ep.Ingress.Denied != nil {
		ep.Ingress.Denied.Sort()
	}
	if ep.Egress.Allowed != nil {
		ep.Egress.Allowed.Sort()
	}
	if ep.Egress.Denied != nil {
		ep.Egress.Denied.Sort()
	}

	return
}

func (e *Endpoint) getEndpointPolicyStateDisabled() (ep *cilium_v2.EndpointPolicy) {
	ep = &cilium_v2.EndpointPolicy{
		Ingress: &cilium_v2.EndpointPolicyDirection{
			State: EndpointPolicyStateDisabled,
		},
		Egress: &cilium_v2.EndpointPolicyDirection{
			State: EndpointPolicyStateDisabled,
		},
	}

	return
}

func (e *Endpoint) getEndpointVisibilityPolicyStatus() *string {
	if e.visibilityPolicy == nil {
		return nil
	}
	var str string
	if e.visibilityPolicy.Error == nil {
		str = "OK"
	} else {
		str = e.visibilityPolicy.Error.Error()
	}
	return &str
}

func (e *Endpoint) getEndpointVisibilityPolicyStatusDisabled() *string {
	str := EndpointPolicyStateDisabled
	return (*string)(&str)
}

// EndpointStatusConfiguration is the configuration interface that a caller of
// to GetCiliumEndpointStatus() must implement
type EndpointStatusConfiguration interface {
	// EndpointStatusIsEnabled must return true if a particular
	// option.EndpointStatus* feature is enabled
	EndpointStatusIsEnabled(option string) bool
}

func compressEndpointState(state models.EndpointState) string {
	switch state {
	case models.EndpointStateRestoring, models.EndpointStateWaitingToRegenerate,
		models.EndpointStateRegenerating, models.EndpointStateReady,
		models.EndpointStateDisconnecting, models.EndpointStateDisconnected:
		return string(models.EndpointStateReady)
	}

	return string(state)
}

// GetCiliumEndpointStatus creates a cilium_v2.EndpointStatus of an endpoint.
// See cilium_v2.EndpointStatus for a detailed explanation of each field.
func (e *Endpoint) GetCiliumEndpointStatus(conf EndpointStatusConfiguration) *cilium_v2.EndpointStatus {
	e.mutex.RLock()
	defer e.mutex.RUnlock()

	status := &cilium_v2.EndpointStatus{
		ID:                  int64(e.ID),
		ExternalIdentifiers: e.getModelEndpointIdentitiersRLocked(),
		Identity:            getEndpointIdentity(identitymodel.CreateModel(e.SecurityIdentity)),
		Networking:          getEndpointNetworking(e.getModelNetworkingRLocked()),
		State:               compressEndpointState(e.getModelCurrentStateRLocked()),
		Encryption:          cilium_v2.EncryptionSpec{Key: int(node.GetIPsecKeyIdentity())},
		NamedPorts:          e.getNamedPortsModel(),
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusControllers) {
		controllerMdl := e.controllers.GetStatusModel()
		status.Controllers = getEndpointStatusControllers(controllerMdl)
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusPolicy) {
		status.Policy = e.getEndpointPolicy()
		status.VisibilityPolicyStatus = e.getEndpointVisibilityPolicyStatus()
	} else {
		status.Policy = e.getEndpointPolicyStateDisabled()
		status.VisibilityPolicyStatus = e.getEndpointVisibilityPolicyStatusDisabled()
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusHealth) {
		status.Health = e.getHealthModel()
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusLog) {
		status.Log = e.getEndpointStatusLog()
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusState) {
		status.State = compressEndpointState(e.getModelCurrentStateRLocked())
	}

	return status
}
<commit_msg>pkg/endpoint: change CEP policy status message<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package endpoint

import (
	"context"
	"sort"
	"strings"
	"time"

	"github.com/cilium/cilium/api/v1/models"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/identity/cache"
	identitymodel "github.com/cilium/cilium/pkg/identity/model"
	cilium_v2 "github.com/cilium/cilium/pkg/k8s/apis/cilium.io/v2"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/node"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/policy"
)

// StatusNA is value of fields in the output of 'kubectl get cep' in case of disabled "--endpoint-status"
const (
	EndpointPolicyStateEnforcing    cilium_v2.EndpointPolicyState = "enforcing"
	EndpointPolicyStateNonEnforcing cilium_v2.EndpointPolicyState = "non-enforcing"
	EndpointPolicyStateDisabled     cilium_v2.EndpointPolicyState = "<status disabled>"
)

func getEndpointStatusControllers(mdlControllers models.ControllerStatuses) (controllers cilium_v2.ControllerList) {
	for _, c := range mdlControllers {
		if c.Status == nil {
			continue
		}

		if c.Status.ConsecutiveFailureCount > 0 {
			s := cilium_v2.ControllerStatus{
				Configuration: c.Configuration,
				Name:          c.Name,
				UUID:          string(c.UUID),
				Status: cilium_v2.ControllerStatusStatus{
					ConsecutiveFailureCount: c.Status.ConsecutiveFailureCount,
					FailureCount:            c.Status.FailureCount,
					LastFailureMsg:          c.Status.LastFailureMsg,
					LastFailureTimestamp:    c.Status.LastFailureTimestamp.String(),
					LastSuccessTimestamp:    c.Status.LastSuccessTimestamp.String(),
					SuccessCount:            c.Status.SuccessCount,
				},
			}
			if controllers == nil {
				controllers = cilium_v2.ControllerList{s}
			} else {
				controllers = append(controllers, s)
			}
		}
	}

	if controllers != nil {
		controllers.Sort()
	}

	return
}

func (e *Endpoint) getEndpointStatusLog() (log []*models.EndpointStatusChange) {
	added := 0

	if s := e.status; s != nil {
		s.indexMU.RLock()
		defer s.indexMU.RUnlock()

		for i := s.lastIndex(); ; i-- {
			if i < 0 {
				i = maxLogs - 1
			}
			if i < len(s.Log) && s.Log[i] != nil {
				l := &models.EndpointStatusChange{
					Timestamp: s.Log[i].Timestamp.Format(time.RFC3339),
					Code:      s.Log[i].Status.Code.String(),
					Message:   s.Log[i].Status.Msg,
					State:     models.EndpointState(s.Log[i].Status.State),
				}

				if strings.ToLower(l.Code) != models.EndpointStatusChangeCodeOk {
					if log == nil {
						log = []*models.EndpointStatusChange{l}
					} else {
						log = append(log, l)
					}

					// Limit the number of endpoint log
					// entries to keep the size of the
					// EndpointStatus low.
					added++
					if added >= cilium_v2.EndpointStatusLogEntries {
						break
					}
				}
			}
			if i == s.Index {
				break
			}
		}
	}
	return
}

func getEndpointIdentity(mdlIdentity *models.Identity) (identity *cilium_v2.EndpointIdentity) {
	if mdlIdentity == nil {
		return
	}
	identity = &cilium_v2.EndpointIdentity{
		ID: mdlIdentity.ID,
	}

	identity.Labels = make([]string, len(mdlIdentity.Labels))
	copy(identity.Labels, mdlIdentity.Labels)
	sort.Strings(identity.Labels)
	return
}

func getEndpointNetworking(mdlNetworking *models.EndpointNetworking) (networking *cilium_v2.EndpointNetworking) {
	if mdlNetworking == nil {
		return nil
	}
	networking = &cilium_v2.EndpointNetworking{
		Addressing: make(cilium_v2.AddressPairList, len(mdlNetworking.Addressing)),
	}

	if option.Config.EnableIPv4 {
		networking.NodeIP = node.GetIPv4().String()
	} else {
		networking.NodeIP = node.GetIPv6().String()
	}

	for i, pair := range mdlNetworking.Addressing {
		networking.Addressing[i] = &cilium_v2.AddressPair{
			IPV4: pair.IPV4,
			IPV6: pair.IPV6,
		}
	}

	networking.Addressing.Sort()
	return
}

// updateLabels inserts the labels correnspoding to the specified identity into
// the AllowedIdentityTuple.
func updateLabels(allocator cache.IdentityAllocator, allowedIdentityTuple *cilium_v2.IdentityTuple, secID identity.NumericIdentity) {
	// IdentityUnknown denotes that this is an L4-only BPF
	// allow, so it applies to all identities. In this case
	// we should skip resolving the labels, because the
	// value 0 does not denote an allow for the "unknown"
	// identity, but instead an allow of all identities for
	// that port.
	if secID != identity.IdentityUnknown {
		identity := allocator.LookupIdentityByID(context.TODO(), secID)
		if identity != nil {
			var l labels.Labels
			if identity.CIDRLabel != nil {
				l = identity.CIDRLabel
			} else {
				l = identity.Labels
			}

			allowedIdentityTuple.IdentityLabels = l.StringMap()
		}
	}
}

// populateResponseWithPolicyKey inserts an AllowedIdentityTuple element into
// 'policy' which corresponds to the specified 'desiredPolicy'. If 'isDeny' is
// true, it will insert the policyKey into the 'Denied'.
func populateResponseWithPolicyKey(
	allocator cache.IdentityAllocator,
	policy *cilium_v2.EndpointPolicy,
	policyKey *policy.Key,
	isDeny bool,
) {
	identityTuple := cilium_v2.IdentityTuple{
		DestPort: policyKey.DestPort,
		Protocol: policyKey.Nexthdr,
		Identity: uint64(policyKey.Identity),
	}

	secID := identity.NumericIdentity(policyKey.Identity)
	updateLabels(allocator, &identityTuple, secID)

	switch {
	case policyKey.IsIngress():
		if isDeny {
			if policy.Ingress.Denied == nil {
				policy.Ingress.Denied = cilium_v2.DenyIdentityList{identityTuple}
			} else {
				policy.Ingress.Denied = append(policy.Ingress.Denied, identityTuple)
			}
		} else {
			if policy.Ingress.Allowed == nil {
				policy.Ingress.Allowed = cilium_v2.AllowedIdentityList{identityTuple}
			} else {
				policy.Ingress.Allowed = append(policy.Ingress.Allowed, identityTuple)
			}
		}
	case policyKey.IsEgress():
		if isDeny {
			if policy.Egress.Denied == nil {
				policy.Egress.Denied = cilium_v2.DenyIdentityList{identityTuple}
			} else {
				policy.Egress.Denied = append(policy.Egress.Denied, identityTuple)
			}
		} else {
			if policy.Egress.Allowed == nil {
				policy.Egress.Allowed = cilium_v2.AllowedIdentityList{identityTuple}
			} else {
				policy.Egress.Allowed = append(policy.Egress.Allowed, identityTuple)
			}
		}
	}
}

// getEndpointPolicy returns an API representation of the policy that the
// received Endpoint intends to apply.
func (e *Endpoint) getEndpointPolicy() (ep *cilium_v2.EndpointPolicy) {
	if e.desiredPolicy == nil {
		return
	}
	ep = &cilium_v2.EndpointPolicy{
		Ingress: &cilium_v2.EndpointPolicyDirection{
			Enforcing: !e.Options.IsEnabled(option.PolicyAuditMode) &&
				e.desiredPolicy.IngressPolicyEnabled,
		},
		Egress: &cilium_v2.EndpointPolicyDirection{
			Enforcing: !e.Options.IsEnabled(option.PolicyAuditMode) &&
				e.desiredPolicy.EgressPolicyEnabled,
		},
	}
	if ep.Ingress.Enforcing {
		ep.Ingress.State = EndpointPolicyStateEnforcing
	} else {
		ep.Ingress.State = EndpointPolicyStateNonEnforcing
	}
	if ep.Egress.Enforcing {
		ep.Egress.State = EndpointPolicyStateEnforcing
	} else {
		ep.Egress.State = EndpointPolicyStateNonEnforcing
	}

	// Handle allow-all cases
	allowsAllIngress, allowsAllEgress := e.desiredPolicy.AllowsIdentity(identity.IdentityUnknown)
	if allowsAllIngress {
		ep.Ingress.Allowed = cilium_v2.AllowedIdentityList{{}}
		ep.Ingress.Denied = cilium_v2.DenyIdentityList{{}}
	}
	if allowsAllEgress {
		ep.Egress.Allowed = cilium_v2.AllowedIdentityList{{}}
		ep.Egress.Denied = cilium_v2.DenyIdentityList{{}}
	}

	// If either ingress or egress policy is enabled, go through
	// the desired policy to populate the values.
	if !allowsAllIngress || !allowsAllEgress {
		allowsWorldIngress, allowsWorldEgress := e.desiredPolicy.AllowsIdentity(identity.ReservedIdentityWorld)

		for policyKey, policyValue := range e.desiredPolicy.PolicyMapState {
			// Skip listing identities if enforcement is disabled in direction,
			// or if the identity corresponds to a CIDR identity and the world is allowed.
			id := identity.NumericIdentity(policyKey.Identity)
			switch {
			case policyKey.IsIngress():
				if allowsAllIngress || (id.HasLocalScope() && allowsWorldIngress) {
					continue
				}
			case policyKey.IsEgress():
				if allowsAllEgress || (id.HasLocalScope() && allowsWorldEgress) {
					continue
				}
			}

			populateResponseWithPolicyKey(e.allocator, ep, &policyKey, policyValue.IsDeny)
		}
	}

	if ep.Ingress.Allowed != nil {
		ep.Ingress.Allowed.Sort()
	}
	if ep.Ingress.Denied != nil {
		ep.Ingress.Denied.Sort()
	}
	if ep.Egress.Allowed != nil {
		ep.Egress.Allowed.Sort()
	}
	if ep.Egress.Denied != nil {
		ep.Egress.Denied.Sort()
	}

	return
}

func (e *Endpoint) getEndpointPolicyStateDisabled() (ep *cilium_v2.EndpointPolicy) {
	ep = &cilium_v2.EndpointPolicy{
		Ingress: &cilium_v2.EndpointPolicyDirection{
			State: EndpointPolicyStateDisabled,
		},
		Egress: &cilium_v2.EndpointPolicyDirection{
			State: EndpointPolicyStateDisabled,
		},
	}

	return
}

func (e *Endpoint) getEndpointVisibilityPolicyStatus() *string {
	if e.visibilityPolicy == nil {
		return nil
	}
	var str string
	if e.visibilityPolicy.Error == nil {
		str = "OK"
	} else {
		str = e.visibilityPolicy.Error.Error()
	}
	return &str
}

func (e *Endpoint) getEndpointVisibilityPolicyStatusDisabled() *string {
	str := EndpointPolicyStateDisabled
	return (*string)(&str)
}

// EndpointStatusConfiguration is the configuration interface that a caller of
// to GetCiliumEndpointStatus() must implement
type EndpointStatusConfiguration interface {
	// EndpointStatusIsEnabled must return true if a particular
	// option.EndpointStatus* feature is enabled
	EndpointStatusIsEnabled(option string) bool
}

func compressEndpointState(state models.EndpointState) string {
	switch state {
	case models.EndpointStateRestoring, models.EndpointStateWaitingToRegenerate,
		models.EndpointStateRegenerating, models.EndpointStateReady,
		models.EndpointStateDisconnecting, models.EndpointStateDisconnected:
		return string(models.EndpointStateReady)
	}

	return string(state)
}

// GetCiliumEndpointStatus creates a cilium_v2.EndpointStatus of an endpoint.
// See cilium_v2.EndpointStatus for a detailed explanation of each field.
func (e *Endpoint) GetCiliumEndpointStatus(conf EndpointStatusConfiguration) *cilium_v2.EndpointStatus {
	e.mutex.RLock()
	defer e.mutex.RUnlock()

	status := &cilium_v2.EndpointStatus{
		ID:                  int64(e.ID),
		ExternalIdentifiers: e.getModelEndpointIdentitiersRLocked(),
		Identity:            getEndpointIdentity(identitymodel.CreateModel(e.SecurityIdentity)),
		Networking:          getEndpointNetworking(e.getModelNetworkingRLocked()),
		State:               compressEndpointState(e.getModelCurrentStateRLocked()),
		Encryption:          cilium_v2.EncryptionSpec{Key: int(node.GetIPsecKeyIdentity())},
		NamedPorts:          e.getNamedPortsModel(),
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusControllers) {
		controllerMdl := e.controllers.GetStatusModel()
		status.Controllers = getEndpointStatusControllers(controllerMdl)
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusPolicy) {
		status.Policy = e.getEndpointPolicy()
		status.VisibilityPolicyStatus = e.getEndpointVisibilityPolicyStatus()
	} else {
		status.Policy = e.getEndpointPolicyStateDisabled()
		status.VisibilityPolicyStatus = e.getEndpointVisibilityPolicyStatusDisabled()
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusHealth) {
		status.Health = e.getHealthModel()
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusLog) {
		status.Log = e.getEndpointStatusLog()
	}

	if conf.EndpointStatusIsEnabled(option.EndpointStatusState) {
		status.State = compressEndpointState(e.getModelCurrentStateRLocked())
	}

	return status
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"flag"
	"io"
	"log"
	"log/syslog"
	"net/http"
	"os"
	"path"
	"time"

	"github.com/tsavola/wag/wasm"
	"golang.org/x/crypto/acme/autocert"

	"github.com/tsavola/gate/run"
	"github.com/tsavola/gate/server"
	_ "github.com/tsavola/gate/service/defaults"
	"github.com/tsavola/gate/service/origin"
)

const (
	renewCertBefore = 30 * 24 * time.Hour

	memorySizeLimit = 256 * wasm.Page
	stackSize       = 16 * 4096
)

func main() {
	var (
		config = run.Config{
			LibDir:   "lib",
			MaxProcs: run.DefaultMaxProcs,
		}
		addr         = "localhost:8888"
		letsencrypt  = false
		email        = ""
		acceptTOS    = false
		certCacheDir = "/var/lib/gate-server-letsencrypt"
		debug        = false
	)

	flag.StringVar(&config.LibDir, "libdir", config.LibDir, "path")
	flag.UintVar(&config.Uids[0], "boot-uid", config.Uids[0], "user id for bootstrapping executor")
	flag.UintVar(&config.Gids[0], "boot-gid", config.Gids[0], "group id for bootstrapping executor")
	flag.UintVar(&config.Uids[1], "exec-uid", config.Uids[1], "user id for executing code")
	flag.UintVar(&config.Gids[1], "exec-gid", config.Gids[1], "group id for executing code")
	flag.UintVar(&config.Gids[2], "pipe-gid", config.Gids[2], "group id for file descriptor sharing")
	flag.StringVar(&config.CgroupParent, "cgroup-parent", config.CgroupParent, "slice")
	flag.StringVar(&config.CgroupTitle, "cgroup-title", config.CgroupTitle, "prefix of dynamic name")
	flag.IntVar(&config.MaxProcs, "max-procs", config.MaxProcs, "limit number of simultaneous programs")
	flag.StringVar(&addr, "addr", addr, "listening [address]:port")
	flag.BoolVar(&letsencrypt, "letsencrypt", letsencrypt, "enable automatic TLS; domain names should be listed after the options")
	flag.StringVar(&email, "email", email, "contact address for Let's Encrypt")
	flag.BoolVar(&acceptTOS, "accept-tos", acceptTOS, "accept Let's Encrypt's terms of service")
	flag.StringVar(&certCacheDir, "cert-cache-dir", certCacheDir, "certificate storage")
	flag.BoolVar(&debug, "debug", debug, "write payload programs' debug output to stderr")
	flag.Parse()
	domains := flag.Args()

	syslogger, err := syslog.New(syslog.LOG_INFO, path.Base(os.Args[0]))
	if err != nil {
		log.Fatal(err)
	}

	env, err := run.NewEnvironment(&config)
	if err != nil {
		syslogger.Crit(err.Error())
		os.Exit(1)
	}
	defer env.Close()

	settings := server.Settings{
		MemorySizeLimit: memorySizeLimit,
		StackSize:       stackSize,
		Env:             env,
		Services:        services,
		Log:             log.New(syslogger, "", 0),
	}

	if debug {
		settings.Debug = os.Stderr
	}

	state := server.NewState(settings)
	handler := server.NewHandler("/", state)

	if letsencrypt {
		if !acceptTOS {
			syslogger.Crit("-accept-tos option not set")
			os.Exit(1)
		}

		m := autocert.Manager{
			Prompt:      autocert.AcceptTOS,
			Cache:       autocert.DirCache(certCacheDir),
			HostPolicy:  autocert.HostWhitelist(domains...),
			RenewBefore: renewCertBefore,
			Email:       email,
		}

		s := http.Server{
			Addr:    addr,
			Handler: handler,
			TLSConfig: &tls.Config{
				GetCertificate: m.GetCertificate,
			},
		}

		err = s.ListenAndServeTLS("", "")
	} else {
		err = http.ListenAndServe(addr, handler)
	}

	syslogger.Crit(err.Error())
	os.Exit(1)
}

func services(r io.Reader, w io.Writer) run.ServiceRegistry {
	return origin.CloneRegistryWith(nil, r, w)
}
<commit_msg>server: use syslog optionally<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"io"
	"log"
	"log/syslog"
	"net/http"
	"os"
	"path"
	"time"

	"github.com/tsavola/wag/wasm"
	"golang.org/x/crypto/acme/autocert"

	"github.com/tsavola/gate/run"
	"github.com/tsavola/gate/server"
	_ "github.com/tsavola/gate/service/defaults"
	"github.com/tsavola/gate/service/origin"
)

const (
	renewCertBefore = 30 * 24 * time.Hour

	memorySizeLimit = 256 * wasm.Page
	stackSize       = 16 * 4096
)

func main() {
	var (
		config = run.Config{
			LibDir:   "lib",
			MaxProcs: run.DefaultMaxProcs,
		}
		addr         = "localhost:8888"
		letsencrypt  = false
		email        = ""
		acceptTOS    = false
		certCacheDir = "/var/lib/gate-server-letsencrypt"
		syslogging   = false
		debug        = false
	)

	flag.StringVar(&config.LibDir, "libdir", config.LibDir, "path")
	flag.UintVar(&config.Uids[0], "boot-uid", config.Uids[0], "user id for bootstrapping executor")
	flag.UintVar(&config.Gids[0], "boot-gid", config.Gids[0], "group id for bootstrapping executor")
	flag.UintVar(&config.Uids[1], "exec-uid", config.Uids[1], "user id for executing code")
	flag.UintVar(&config.Gids[1], "exec-gid", config.Gids[1], "group id for executing code")
	flag.UintVar(&config.Gids[2], "pipe-gid", config.Gids[2], "group id for file descriptor sharing")
	flag.StringVar(&config.CgroupParent, "cgroup-parent", config.CgroupParent, "slice")
	flag.StringVar(&config.CgroupTitle, "cgroup-title", config.CgroupTitle, "prefix of dynamic name")
	flag.IntVar(&config.MaxProcs, "max-procs", config.MaxProcs, "limit number of simultaneous programs")
	flag.StringVar(&addr, "addr", addr, "listening [address]:port")
	flag.BoolVar(&letsencrypt, "letsencrypt", letsencrypt, "enable automatic TLS; domain names should be listed after the options")
	flag.StringVar(&email, "email", email, "contact address for Let's Encrypt")
	flag.BoolVar(&acceptTOS, "accept-tos", acceptTOS, "accept Let's Encrypt's terms of service")
	flag.StringVar(&certCacheDir, "cert-cache-dir", certCacheDir, "certificate storage")
	flag.BoolVar(&syslogging, "syslog", syslogging, "send log messages to syslog instead of stderr")
	flag.BoolVar(&debug, "debug", debug, "write payload programs' debug output to stderr")
	flag.Parse()
	domains := flag.Args()

	var (
		critLog *log.Logger
		infoLog server.Logger
	)

	if syslogging {
		tag := path.Base(os.Args[0])

		w, err := syslog.New(syslog.LOG_CRIT, tag)
		if err != nil {
			log.Fatal(err)
		}
		critLog = log.New(w, "", 0)

		w, err = syslog.New(syslog.LOG_INFO, tag)
		if err != nil {
			critLog.Fatal(err)
		}
		infoLog = log.New(w, "", 0)
	} else {
		critLog = log.New(os.Stderr, "", 0)
		infoLog = critLog
	}

	env, err := run.NewEnvironment(&config)
	if err != nil {
		critLog.Fatal(err)
	}
	defer env.Close()

	settings := server.Settings{
		MemorySizeLimit: memorySizeLimit,
		StackSize:       stackSize,
		Env:             env,
		Services:        services,
		Log:             infoLog,
	}

	if debug {
		settings.Debug = os.Stderr
	}

	state := server.NewState(settings)
	handler := server.NewHandler("/", state)

	if letsencrypt {
		if !acceptTOS {
			critLog.Fatal("-accept-tos option not set")
		}

		m := autocert.Manager{
			Prompt:      autocert.AcceptTOS,
			Cache:       autocert.DirCache(certCacheDir),
			HostPolicy:  autocert.HostWhitelist(domains...),
			RenewBefore: renewCertBefore,
			Email:       email,
		}

		s := http.Server{
			Addr:    addr,
			Handler: handler,
			TLSConfig: &tls.Config{
				GetCertificate: m.GetCertificate,
			},
		}

		err = s.ListenAndServeTLS("", "")
	} else {
		err = http.ListenAndServe(addr, handler)
	}

	critLog.Fatal(err)
}

func services(r io.Reader, w io.Writer) run.ServiceRegistry {
	return origin.CloneRegistryWith(nil, r, w)
}
<|endoftext|>


<commit_before>package templates

import (
	"bytes"
	"net/http"
	"path"
	"text/template"

	"github.com/mholt/caddy/middleware"
)

// New constructs a new Templates middleware instance.
func New(c middleware.Controller) (middleware.Middleware, error) {
	rules, err := parse(c)
	if err != nil {
		return nil, err
	}

	tmpls := Templates{
		Root:  c.Root(),
		Rules: rules,
	}

	return func(next middleware.Handler) middleware.Handler {
		tmpls.Next = next
		return tmpls
	}, nil
}

// ServeHTTP implements the middleware.Handler interface.
func (t Templates) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
	for _, rule := range t.Rules {
		if !middleware.Path(r.URL.Path).Matches(rule.Path) {
			continue
		}

		reqExt := path.Ext(r.URL.Path)

		for _, ext := range rule.Extensions {
			if reqExt == ext {
				// Create execution context
				ctx := context{root: http.Dir(t.Root), req: r, URL: r.URL}

				// Build the template
				tpl, err := template.ParseFiles(t.Root + r.URL.Path)
				if err != nil {
					return http.StatusInternalServerError, err
				}

				// Execute it
				var buf bytes.Buffer
				err = bc.Template.Execute(&buf, ctx)
				if err != nil {
					return http.StatusInternalServerError, err
				}
				buf.WriteTo(w)

				return http.StatusOK, nil
			}
		}
	}

	return t.Next.ServeHTTP(w, r)
}

func parse(c middleware.Controller) ([]Rule, error) {
	var rules []Rule

	for c.Next() {
		var rule Rule

		if c.NextArg() {
			// First argument would be the path
			rule.Path = c.Val()

			// Any remaining arguments are extensions
			rule.Extensions = c.RemainingArgs()
			if len(rule.Extensions) == 0 {
				rule.Extensions = defaultExtensions
			}
		} else {
			rule.Path = defaultPath
			rule.Extensions = defaultExtensions
		}

		rules = append(rules, rule)
	}

	return rules, nil
}

// Templates is middleware to render templated files as the HTTP response.
type Templates struct {
	Next  middleware.Handler
	Root  string
	Rules []Rule
}

// Rule represents a template rule. A template will only execute
// with this rule if the request path matches the Path specified
// and requests a resource with one of the extensions specified.
type Rule struct {
	Path       string
	Extensions []string
}

const defaultPath = "/"

var defaultExtensions = []string{".html", ".htm", ".txt"}
<commit_msg>Make the thing compile<commit_after>package templates

import (
	"bytes"
	"net/http"
	"path"
	"text/template"

	"github.com/mholt/caddy/middleware"
)

// New constructs a new Templates middleware instance.
func New(c middleware.Controller) (middleware.Middleware, error) {
	rules, err := parse(c)
	if err != nil {
		return nil, err
	}

	tmpls := Templates{
		Root:  c.Root(),
		Rules: rules,
	}

	return func(next middleware.Handler) middleware.Handler {
		tmpls.Next = next
		return tmpls
	}, nil
}

// ServeHTTP implements the middleware.Handler interface.
func (t Templates) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
	for _, rule := range t.Rules {
		if !middleware.Path(r.URL.Path).Matches(rule.Path) {
			continue
		}

		reqExt := path.Ext(r.URL.Path)

		for _, ext := range rule.Extensions {
			if reqExt == ext {
				// Create execution context
				ctx := context{root: http.Dir(t.Root), req: r, URL: r.URL}

				// Build the template
				tpl, err := template.ParseFiles(t.Root + r.URL.Path)
				if err != nil {
					return http.StatusInternalServerError, err
				}

				// Execute it
				var buf bytes.Buffer
				err = tpl.Execute(&buf, ctx)
				if err != nil {
					return http.StatusInternalServerError, err
				}
				buf.WriteTo(w)

				return http.StatusOK, nil
			}
		}
	}

	return t.Next.ServeHTTP(w, r)
}

func parse(c middleware.Controller) ([]Rule, error) {
	var rules []Rule

	for c.Next() {
		var rule Rule

		if c.NextArg() {
			// First argument would be the path
			rule.Path = c.Val()

			// Any remaining arguments are extensions
			rule.Extensions = c.RemainingArgs()
			if len(rule.Extensions) == 0 {
				rule.Extensions = defaultExtensions
			}
		} else {
			rule.Path = defaultPath
			rule.Extensions = defaultExtensions
		}

		rules = append(rules, rule)
	}

	return rules, nil
}

// Templates is middleware to render templated files as the HTTP response.
type Templates struct {
	Next  middleware.Handler
	Root  string
	Rules []Rule
}

// Rule represents a template rule. A template will only execute
// with this rule if the request path matches the Path specified
// and requests a resource with one of the extensions specified.
type Rule struct {
	Path       string
	Extensions []string
}

const defaultPath = "/"

var defaultExtensions = []string{".html", ".htm", ".txt"}
<|endoftext|>


<commit_before>package markdown

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/mholt/caddy/middleware"
)

// GenerateStatic generate static files and link index from markdowns.
// It only generates static files if it is enabled (cfg.StaticDir
// must be set).
func GenerateStatic(md Markdown, cfg *Config) error {
	generated, err := generateLinks(md, cfg)
	if err != nil {
		return err
	}

	// No new file changes, return.
	if !generated {
		return nil
	}

	// If static site generation is enabled.
	if cfg.StaticDir != "" {
		if err := generateStaticHTML(md, cfg); err != nil {
			return err
		}
	}
	return nil
}

type linkGenerator struct {
	gens map[*Config]*linkGen
	sync.Mutex
}

var generator = linkGenerator{gens: make(map[*Config]*linkGen)}

// generateLinks generates links to all markdown files ordered by newest date.
// This blocks until link generation is done. When called by multiple goroutines,
// the first caller starts the generation and others only wait.
// It returns if generation is done and any error that occurred.
func generateLinks(md Markdown, cfg *Config) (bool, error) {
	generator.Lock()

	// if link generator exists for config and running, wait.
	if g, ok := generator.gens[cfg]; ok {
		if g.started() {
			g.addWaiter()
			generator.Unlock()
			g.Wait()
			// another goroutine has done the generation.
			return false, g.lastErr
		}
	}

	g := &linkGen{}
	generator.gens[cfg] = g
	generator.Unlock()

	generated := g.generateLinks(md, cfg)
	g.discardWaiters()
	return generated, g.lastErr
}

// generateStaticHTML generates static HTML files from markdowns.
func generateStaticHTML(md Markdown, cfg *Config) error {
	// If generated site already exists, clear it out
	_, err := os.Stat(cfg.StaticDir)
	if err == nil {
		err := os.RemoveAll(cfg.StaticDir)
		if err != nil {
			return err
		}
	}

	fp := filepath.Join(md.Root, cfg.PathScope)

	return filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
		for _, ext := range cfg.Extensions {
			if !info.IsDir() && strings.HasSuffix(info.Name(), ext) {
				// Load the file
				body, err := ioutil.ReadFile(path)
				if err != nil {
					return err
				}

				// Get the relative path as if it were a HTTP request,
				// then prepend with "/" (like a real HTTP request)
				reqPath, err := filepath.Rel(md.Root, path)
				if err != nil {
					return err
				}
				reqPath = filepath.ToSlash(reqPath)
				reqPath = "/" + reqPath

				// Generate the static file
				ctx := middleware.Context{Root: md.FileSys}
				_, err = md.Process(cfg, reqPath, body, ctx)
				if err != nil {
					return err
				}

				break // don't try other file extensions
			}
		}
		return nil
	})
}

// computeDirHash computes an hash on static directory of c.
func computeDirHash(md Markdown, c *Config) (string, error) {
	dir := filepath.Join(md.Root, c.PathScope)
	if _, err := os.Stat(dir); err != nil {
		return "", err
	}

	hashString := ""
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() && c.IsValidExt(filepath.Ext(path)) {
			hashString += fmt.Sprintf("%v%v%v%v", info.ModTime(), info.Name(), info.Size(), path)
		}
		return nil
	})
	if err != nil {
		return "", err
	}

	sum := md5.Sum([]byte(hashString))
	return hex.EncodeToString(sum[:]), nil
}
<commit_msg>markdown: Don't generate static site or links unless sitegen is enabled<commit_after>package markdown

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/mholt/caddy/middleware"
)

// GenerateStatic generate static files and link index from markdowns.
// It only generates static files if it is enabled (cfg.StaticDir
// must be set).
func GenerateStatic(md Markdown, cfg *Config) error {
	// If static site generation is enabled.
	if cfg.StaticDir != "" {
		generated, err := generateLinks(md, cfg)
		if err != nil {
			return err
		}

		// No new file changes, return.
		if !generated {
			return nil
		}

		if err := generateStaticHTML(md, cfg); err != nil {
			return err
		}
	}
	return nil
}

type linkGenerator struct {
	gens map[*Config]*linkGen
	sync.Mutex
}

var generator = linkGenerator{gens: make(map[*Config]*linkGen)}

// generateLinks generates links to all markdown files ordered by newest date.
// This blocks until link generation is done. When called by multiple goroutines,
// the first caller starts the generation and others only wait.
// It returns if generation is done and any error that occurred.
func generateLinks(md Markdown, cfg *Config) (bool, error) {
	generator.Lock()

	// if link generator exists for config and running, wait.
	if g, ok := generator.gens[cfg]; ok {
		if g.started() {
			g.addWaiter()
			generator.Unlock()
			g.Wait()
			// another goroutine has done the generation.
			return false, g.lastErr
		}
	}

	g := &linkGen{}
	generator.gens[cfg] = g
	generator.Unlock()

	generated := g.generateLinks(md, cfg)
	g.discardWaiters()
	return generated, g.lastErr
}

// generateStaticHTML generates static HTML files from markdowns.
func generateStaticHTML(md Markdown, cfg *Config) error {
	// If generated site already exists, clear it out
	_, err := os.Stat(cfg.StaticDir)
	if err == nil {
		err := os.RemoveAll(cfg.StaticDir)
		if err != nil {
			return err
		}
	}

	fp := filepath.Join(md.Root, cfg.PathScope)

	return filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
		for _, ext := range cfg.Extensions {
			if !info.IsDir() && strings.HasSuffix(info.Name(), ext) {
				// Load the file
				body, err := ioutil.ReadFile(path)
				if err != nil {
					return err
				}

				// Get the relative path as if it were a HTTP request,
				// then prepend with "/" (like a real HTTP request)
				reqPath, err := filepath.Rel(md.Root, path)
				if err != nil {
					return err
				}
				reqPath = filepath.ToSlash(reqPath)
				reqPath = "/" + reqPath

				// Generate the static file
				ctx := middleware.Context{Root: md.FileSys}
				_, err = md.Process(cfg, reqPath, body, ctx)
				if err != nil {
					return err
				}

				break // don't try other file extensions
			}
		}
		return nil
	})
}

// computeDirHash computes an hash on static directory of c.
func computeDirHash(md Markdown, c *Config) (string, error) {
	dir := filepath.Join(md.Root, c.PathScope)
	if _, err := os.Stat(dir); err != nil {
		return "", err
	}

	hashString := ""
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() && c.IsValidExt(filepath.Ext(path)) {
			hashString += fmt.Sprintf("%v%v%v%v", info.ModTime(), info.Name(), info.Size(), path)
		}
		return nil
	})
	if err != nil {
		return "", err
	}

	sum := md5.Sum([]byte(hashString))
	return hex.EncodeToString(sum[:]), nil
}
<|endoftext|>


<commit_before>// Copyright © 2015 Steve Francia <spf@spf13.com>.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var cmdDirs = [...]string{"cmd", "cmds", "command", "commands"}
var srcPaths []string

func init() {
	// Initialize srcPaths.
	envGoPath := os.Getenv("GOPATH")
	goPaths := filepath.SplitList(envGoPath)
	if len(goPaths) == 0 {
		er("$GOPATH is not set")
	}
	srcPaths = make([]string, 0, len(goPaths))
	for _, goPath := range goPaths {
		srcPaths = append(srcPaths, filepath.Join(goPath, "src"))
	}
}

func er(msg interface{}) {
	fmt.Println("Error:", msg)
	os.Exit(1)
}

// isEmpty checks if a given path is empty.
func isEmpty(path string) bool {
	fi, err := os.Stat(path)
	if err != nil {
		er(err)
	}
	if fi.IsDir() {
		f, err := os.Open(path)
		if err != nil {
			er(err)
		}
		defer f.Close()
		dirs, err := f.Readdirnames(1)
		if err != nil && err != io.EOF {
			er(err)
		}
		return len(dirs) == 0
	}
	return fi.Size() == 0
}

// exists checks if a file or directory exists.
func exists(path string) bool {
	if path == "" {
		return false
	}
	_, err := os.Stat(path)
	if err == nil {
		return true
	}
	if !os.IsNotExist(err) {
		er(err)
	}
	return false
}

func executeTemplate(tmplStr string, data interface{}) (string, error) {
	tmpl, err := template.New("").Funcs(template.FuncMap{"comment": commentifyString}).Parse(tmplStr)
	if err != nil {
		return "", err
	}

	buf := new(bytes.Buffer)
	err = tmpl.Execute(buf, data)
	return buf.String(), err
}

func writeStringToFile(path string, s string) error {
	return writeToFile(path, strings.NewReader(s))
}

// writeToFile writes r to file with path only
// if file/directory on given path doesn't exist.
// If file/directory exists on given path, then
// it terminates app and prints an appropriate error.
func writeToFile(path string, r io.Reader) error {
	if exists(path) {
		return fmt.Errorf("%v already exists", path)
	}

	dir := filepath.Dir(path)
	if dir != "" {
		if err := os.MkdirAll(dir, 0777); err != nil {
			return err
		}
	}

	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, r)
	return err
}

// commentfyString comments every line of in.
func commentifyString(in string) string {
	var newlines []string
	lines := strings.Split(in, "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "//") {
			newlines = append(newlines, line)
		} else {
			if line == "" {
				newlines = append(newlines, "//")
			} else {
				newlines = append(newlines, "// "+line)
			}
		}
	}
	return strings.Join(newlines, "\n")
}
<commit_msg>cmd: Ignore hidden files in isEmpty<commit_after>// Copyright © 2015 Steve Francia <spf@spf13.com>.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var cmdDirs = [...]string{"cmd", "cmds", "command", "commands"}
var srcPaths []string

func init() {
	// Initialize srcPaths.
	envGoPath := os.Getenv("GOPATH")
	goPaths := filepath.SplitList(envGoPath)
	if len(goPaths) == 0 {
		er("$GOPATH is not set")
	}
	srcPaths = make([]string, 0, len(goPaths))
	for _, goPath := range goPaths {
		srcPaths = append(srcPaths, filepath.Join(goPath, "src"))
	}
}

func er(msg interface{}) {
	fmt.Println("Error:", msg)
	os.Exit(1)
}

// isEmpty checks if a given path is empty.
// Hidden files in path are ignored.
func isEmpty(path string) bool {
	fi, err := os.Stat(path)
	if err != nil {
		er(err)
	}

	if !fi.IsDir() {
		return fi.Size() == 0
	}

	f, err := os.Open(path)
	if err != nil {
		er(err)
	}
	defer f.Close()

	names, err := f.Readdirnames(-1)
	if err != nil && err != io.EOF {
		er(err)
	}

	for _, name := range names {
		if len(name) > 0 && name[0] != '.' {
			return false
		}
	}
	return true
}

// exists checks if a file or directory exists.
func exists(path string) bool {
	if path == "" {
		return false
	}
	_, err := os.Stat(path)
	if err == nil {
		return true
	}
	if !os.IsNotExist(err) {
		er(err)
	}
	return false
}

func executeTemplate(tmplStr string, data interface{}) (string, error) {
	tmpl, err := template.New("").Funcs(template.FuncMap{"comment": commentifyString}).Parse(tmplStr)
	if err != nil {
		return "", err
	}

	buf := new(bytes.Buffer)
	err = tmpl.Execute(buf, data)
	return buf.String(), err
}

func writeStringToFile(path string, s string) error {
	return writeToFile(path, strings.NewReader(s))
}

// writeToFile writes r to file with path only
// if file/directory on given path doesn't exist.
// If file/directory exists on given path, then
// it terminates app and prints an appropriate error.
func writeToFile(path string, r io.Reader) error {
	if exists(path) {
		return fmt.Errorf("%v already exists", path)
	}

	dir := filepath.Dir(path)
	if dir != "" {
		if err := os.MkdirAll(dir, 0777); err != nil {
			return err
		}
	}

	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, r)
	return err
}

// commentfyString comments every line of in.
func commentifyString(in string) string {
	var newlines []string
	lines := strings.Split(in, "\n")
	for _, line := range lines {
		if strings.HasPrefix(line, "//") {
			newlines = append(newlines, line)
		} else {
			if line == "" {
				newlines = append(newlines, "//")
			} else {
				newlines = append(newlines, "// "+line)
			}
		}
	}
	return strings.Join(newlines, "\n")
}
<|endoftext|>


<commit_before>package testcover

import (
	"encoding/binary"
)

func Fuzz(data []byte) int {
	if len(data) < 20 {
		return 0
	}
	x := binary.BigEndian.Uint32(data[12:])
	if x == 0x45839281 {
		bingo()
	}
	if data[10] == 0xfd && data[15] == 0x9a && data[17] == 0x71 {
		bingo()
	}
	switch binary.LittleEndian.Uint32(data[10:]) {
	default:
		bingo()
	case 0x12345678:
		bingo()
	case 0x98765432:
		bingo()
	}
	switch {
	case binary.LittleEndian.Uint32(data[8:]) == 0x12345678:
		bingo()
	default:
		bingo()
	case 0x98765432 == binary.BigEndian.Uint32(data[7:]):
		bingo()
	}

	switch string(data[5:9]) {
	case "ABCD":
		bingo()
	case "QWER":
		bingo()
	case "ZXCV":
		bingo()
	}

	n := binary.BigEndian.Uint32(data[0:4])
	if int(n) <= len(data)-4 {
		s := string(data[4 : 4+n])
		if s == "eat this" {
			bingo()
		}
	}

	if f := binary.BigEndian.Uint32(data[9:]) > 0xfffffffd; f {
		bingo()
	}

	if len(data) > 20 {
		magic := uint32(data[5]) | uint32(data[6])<<8 | uint32(data[9])<<16 | uint32(data[11])<<24
		if magic == 0xabcd1234 {
			bingo()
		}
	}

	return 0
}

func bingo() {
	if theFalse {
		bingo()
	}
}

var theFalse = false
<commit_msg>support (1) custom string types and (2) byte arrays in sonar<commit_after>package testcover

import (
	"bytes"
	"encoding/binary"
)

func Fuzz(data []byte) int {
	if len(data) < 20 {
		return 0
	}
	x := binary.BigEndian.Uint32(data[12:])
	if x == 0x45839281 {
		bingo()
	}
	if data[10] == 0xfd && data[15] == 0x9a && data[17] == 0x71 {
		bingo()
	}
	switch binary.LittleEndian.Uint32(data[10:]) {
	default:
		bingo()
	case 0x12345678:
		bingo()
	case 0x98765432:
		bingo()
	}
	switch {
	case binary.LittleEndian.Uint32(data[8:]) == 0x12345678:
		bingo()
	default:
		bingo()
	case 0x98765432 == binary.BigEndian.Uint32(data[7:]):
		bingo()
	}

	switch string(data[5:9]) {
	case "ABCD":
		bingo()
	case "QWER":
		bingo()
	case "ZXCV":
		bingo()
	}

	n := binary.BigEndian.Uint32(data[0:4])
	if int(n) <= len(data)-4 {
		s := string(data[4 : 4+n])
		if s == "eat this" {
			bingo()
		}
	}

	if f := binary.BigEndian.Uint32(data[9:]) > 0xfffffffd; f {
		bingo()
	}

	magic := uint32(data[5]) | uint32(data[6])<<8 | uint32(data[9])<<16 | uint32(data[11])<<24
	if magic == 0xabcd1234 {
		bingo()
	}

	type Hdr struct {
		Magic [8]byte
		N     uint32
	}
	var hdr Hdr
	binary.Read(bytes.NewReader(data), binary.LittleEndian, &hdr)
	if hdr.Magic == [8]byte{'m', 'a', 'g', 'i', 'c', 'h', 'd', 'r'} {
		bingo()
	}

	type Name string
	name := Name(data[4:9])
	if name == "12345" {
		bingo()
	}

	return 0
}

func bingo() {
	if theFalse {
		bingo()
	}
}

var theFalse = false
<|endoftext|>


<commit_before>package msg

import (
	"cloud.google.com/go/pubsub"
	"context"
	"fmt"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/cred"
	context2 "golang.org/x/net/context"
	"google.golang.org/api/option"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

type gcPubSubClient struct {
	ctx       context.Context
	client    *pubsub.Client
	projectId string
	timeout   time.Duration
}

func (s *gcPubSubClient) createSubscription(resource *ResourceSetup) (*pubsub.Subscription, error) {
	subscription, err := s.getSubscription(&resource.Resource)
	if err != nil {
		return nil, err
	}
	exists, err := subscription.Exists(s.ctx)
	if err != nil {
		return nil, err
	}
	if exists && resource.Recreate {
		exists = false
		if err = subscription.Delete(s.ctx); err != nil {
			return nil, err
		}
	}
	if !exists {
		config := resource.Config
		if config.Topic == nil {
			return nil, fmt.Errorf("topic was empty")
		}
		topic, err := s.getTopic(config.Topic)
		if err != nil {
			return nil, err
		}

		subscriptionConfig := pubsub.SubscriptionConfig{
			Labels:              config.Labels,
			Topic:               topic,
			AckDeadline:         config.AckDeadline,
			RetainAckedMessages: config.RetainAckedMessages,
			RetentionDuration:   config.RetentionDuration,
		}

		subscription, err = s.client.CreateSubscription(s.ctx, subscription.ID(), subscriptionConfig)
	}
	return subscription, err
}

func (s *gcPubSubClient) createTopic(resource *ResourceSetup) (*pubsub.Topic, error) {
	topic, err := s.getTopic(&resource.Resource)
	if err != nil {
		return nil, err
	}
	exists, err := topic.Exists(s.ctx)
	if err != nil {
		return nil, err
	}
	if exists && resource.Recreate {
		exists = false
		if err = topic.Delete(s.ctx); err != nil {
			return nil, err
		}
	}

	if !exists {
		if topic, err = s.client.CreateTopic(s.ctx, topic.ID()); err != nil {
			return nil, err
		}
	}
	return topic, nil
}

func (s *gcPubSubClient) Create(resource *ResourceSetup) (*Resource, error) {
	var err error
	var result = resource.Resource
	switch resource.Type {
	case ResourceTypeTopic:
		topic, err := s.createTopic(resource)
		if err != nil {
			return nil, err
		}
		result.ID = topic.ID()
	case ResourceTypeSubscription:
		subscription, err := s.createSubscription(resource)
		if err != nil {
			return nil, err
		}
		result.ID = subscription.ID()
	default:
		err = fmt.Errorf("unsupported resource type: %v, %v", resource.Type, resource.URL)
	}
	return &result, err
}

func (s *gcPubSubClient) Delete(resource *Resource) error {
	switch resource.Type {
	case ResourceTypeTopic:
		topic, err := s.getTopic(resource)
		if err != nil {
			return err
		}
		return topic.Delete(s.ctx)
	case ResourceTypeSubscription:
		subscription, err := s.getSubscription(resource)
		if err != nil {
			return err
		}
		return subscription.Delete(s.ctx)
	}
	return fmt.Errorf("unsupported resource type: %v, %v", resource.Type, resource.URL)
}

func (s *gcPubSubClient) Push(dest *Resource, message *Message) (Result, error) {
	topic, err := s.getTopic(dest)
	if err != nil {
		return nil, err
	}
	ok, err := topic.Exists(s.ctx)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("topic %v does not exist", dest)
	}
	var pubMessage = &pubsub.Message{
		Attributes: message.Attributes,
	}

	if message.Data != nil {
		switch data := message.Data.(type) {
		case string:
			pubMessage.Data = []byte(data)
		case []byte:
			pubMessage.Data = data
		default:
			JSONText, err := toolbox.AsIndentJSONText(data)
			if err != nil {
				return nil, err
			}
			pubMessage.Data = []byte(JSONText)
		}
	}

	response := topic.Publish(s.ctx, pubMessage)
	serverId, err := response.Get(s.ctx)
	if err != nil {
		return response, err
	}
	select {
	case <-response.Ready():
	case <-time.After(s.timeout):
		log.Printf("publish ready timeout reached: %s", s.timeout)
	}
	return serverId, err
}

func (s *gcPubSubClient) PullN(source *Resource, max int, nack bool) ([]*Message, error) {
	subscription, err := s.getSubscription(source)
	if err != nil {
		return nil, err
	}
	var pulledCounter int32 = 0
	ctx, cancel := context.WithCancel(s.ctx)
	go func() {
		time.Sleep(s.timeout)

		pulledCount := int(atomic.AddInt32(&pulledCounter, 1))
		if pulledCount != max && max > 0 {
			log.Printf("publish ready timeout reached: %s", s.timeout)
			cancel()
		}
	}()
	mutex := &sync.Mutex{}
	var messages = make([]*Message, 0)
	err = subscription.Receive(ctx, func(ctx context2.Context, msg *pubsub.Message) {
		mutex.Lock()
		defer mutex.Unlock()
		messages = append(messages, &Message{
			ID:         msg.ID,
			Attributes: msg.Attributes,
			Data:       msg.Data,
		})
		pulledCount := int(atomic.AddInt32(&pulledCounter, 1))
		if max > 0 && pulledCount >= max {
			cancel()
		}
		if nack {
			msg.Nack()
		} else {
			msg.Ack()
		}
	})
	return messages, err
}

func (s *gcPubSubClient) Close() error {
	return s.client.Close()
}

func (s *gcPubSubClient) getSubscription(dest *Resource) (*pubsub.Subscription, error) {
	if dest.Name == "" && dest.URL == "" {
		return nil, fmt.Errorf("subscription name and URL was empty, expected /projects/[PROJECT ID]/subscriptions/[SUBSCRIPTION] URL or subscription name")
	}
	if dest.projectID == "" {
		return s.client.Subscription(dest.Name), nil
	}
	return s.client.SubscriptionInProject(dest.Name, dest.projectID), nil
}

func (s *gcPubSubClient) getTopic(dest *Resource) (*pubsub.Topic, error) {
	if dest.Name == "" && dest.URL == "" {
		return nil, fmt.Errorf("subscription name and URL was empty, expected /projects/[PROJECT ID]/topics/[SUBSCRIPTION] URL or topic name")
	}
	if dest.projectID == "" {
		return s.client.Topic(dest.Name), nil
	}
	return s.client.TopicInProject(dest.Name, dest.projectID), nil
}

func newCloudPubSub(credConfig *cred.Config, URL string, timeout time.Duration) (Client, error) {
	ctx := context.Background()
	jwtConfig, err := credConfig.NewJWTConfig(pubsub.ScopePubSub)
	if err != nil {
		return nil, err
	}
	var opts = []option.ClientOption{
		option.WithTokenSource(jwtConfig.TokenSource(ctx)),
	}
	client, err := pubsub.NewClient(ctx, credConfig.ProjectID, opts...)
	if err != nil {
		return nil, fmt.Errorf("failed to create pubsub client: %v", err)
	}
	var projectID = extractSubPath(URL, "project")
	if projectID == "" {
		projectID = credConfig.ProjectID
	}

	var service = &gcPubSubClient{
		client:    client,
		ctx:       ctx,
		timeout:   timeout,
		projectId: projectID,
	}
	return service, nil
}
<commit_msg>updated error handling<commit_after>package msg

import (
	"cloud.google.com/go/pubsub"
	"context"
	"fmt"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/cred"
	context2 "golang.org/x/net/context"
	"google.golang.org/api/option"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

type gcPubSubClient struct {
	ctx       context.Context
	client    *pubsub.Client
	projectId string
	timeout   time.Duration
}

func (s *gcPubSubClient) createSubscription(resource *ResourceSetup) (*pubsub.Subscription, error) {
	subscription, err := s.getSubscription(&resource.Resource)
	if err != nil {
		return nil, err
	}
	exists, err := subscription.Exists(s.ctx)
	if err != nil {
		return nil, err
	}

	if exists && resource.Recreate {
		exists = false
		if err = subscription.Delete(s.ctx); err != nil {
			return nil, err
		}
	}

	if !exists {
		config := resource.Config
		if config.Topic == nil {
			return nil, fmt.Errorf("topic was empty")
		}
		topic, err := s.getTopic(config.Topic)
		if err != nil {
			return nil, err
		}

		subscriptionConfig := pubsub.SubscriptionConfig{
			Labels:              config.Labels,
			Topic:               topic,
			AckDeadline:         config.AckDeadline,
			RetainAckedMessages: config.RetainAckedMessages,
			RetentionDuration:   config.RetentionDuration,
		}

		if subscription, err = s.client.CreateSubscription(s.ctx, subscription.ID(), subscriptionConfig); err != nil {
			return nil, err
		}
	}
	return subscription, err
}

func (s *gcPubSubClient) createTopic(resource *ResourceSetup) (*pubsub.Topic, error) {
	topic, err := s.getTopic(&resource.Resource)
	if err != nil {
		return nil, err
	}
	exists, err := topic.Exists(s.ctx)
	if err != nil {
		return nil, err
	}
	if exists && resource.Recreate {
		exists = false
		if err = topic.Delete(s.ctx); err != nil {
			return nil, err
		}
	}

	if !exists {
		if topic, err = s.client.CreateTopic(s.ctx, topic.ID()); err != nil {
			return nil, err
		}
	}
	return topic, nil
}

func (s *gcPubSubClient) Create(resource *ResourceSetup) (*Resource, error) {
	var err error
	var result = resource.Resource
	switch resource.Type {
	case ResourceTypeTopic:
		topic, err := s.createTopic(resource)
		if err != nil {
			return nil, err
		}
		result.ID = topic.ID()
	case ResourceTypeSubscription:
		subscription, err := s.createSubscription(resource)
		if err != nil {
			return nil, err
		}
		result.ID = subscription.ID()
	default:
		err = fmt.Errorf("unsupported resource type: %v, %v", resource.Type, resource.URL)
	}
	return &result, err
}

func (s *gcPubSubClient) Delete(resource *Resource) error {
	switch resource.Type {
	case ResourceTypeTopic:
		topic, err := s.getTopic(resource)
		if err != nil {
			return err
		}
		return topic.Delete(s.ctx)
	case ResourceTypeSubscription:
		subscription, err := s.getSubscription(resource)
		if err != nil {
			return err
		}
		return subscription.Delete(s.ctx)
	}
	return fmt.Errorf("unsupported resource type: %v, %v", resource.Type, resource.URL)
}

func (s *gcPubSubClient) Push(dest *Resource, message *Message) (Result, error) {
	topic, err := s.getTopic(dest)
	if err != nil {
		return nil, err
	}
	ok, err := topic.Exists(s.ctx)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("topic %v does not exist", dest)
	}
	var pubMessage = &pubsub.Message{
		Attributes: message.Attributes,
	}

	if message.Data != nil {
		switch data := message.Data.(type) {
		case string:
			pubMessage.Data = []byte(data)
		case []byte:
			pubMessage.Data = data
		default:
			JSONText, err := toolbox.AsIndentJSONText(data)
			if err != nil {
				return nil, err
			}
			pubMessage.Data = []byte(JSONText)
		}
	}

	response := topic.Publish(s.ctx, pubMessage)
	serverId, err := response.Get(s.ctx)
	if err != nil {
		return response, err
	}
	select {
	case <-response.Ready():
	case <-time.After(s.timeout):
		log.Printf("publish ready timeout reached: %s", s.timeout)
	}
	return serverId, err
}

func (s *gcPubSubClient) PullN(source *Resource, max int, nack bool) ([]*Message, error) {
	subscription, err := s.getSubscription(source)
	if err != nil {
		return nil, err
	}
	var pulledCounter int32 = 0
	ctx, cancel := context.WithCancel(s.ctx)
	go func() {
		time.Sleep(s.timeout)

		pulledCount := int(atomic.AddInt32(&pulledCounter, 1))
		if pulledCount != max && max > 0 {
			log.Printf("publish ready timeout reached: %s", s.timeout)
			cancel()
		}
	}()
	mutex := &sync.Mutex{}
	var messages = make([]*Message, 0)
	err = subscription.Receive(ctx, func(ctx context2.Context, msg *pubsub.Message) {
		mutex.Lock()
		defer mutex.Unlock()
		messages = append(messages, &Message{
			ID:         msg.ID,
			Attributes: msg.Attributes,
			Data:       msg.Data,
		})
		pulledCount := int(atomic.AddInt32(&pulledCounter, 1))
		if max > 0 && pulledCount >= max {
			cancel()
		}
		if nack {
			msg.Nack()
		} else {
			msg.Ack()
		}
	})
	return messages, err
}

func (s *gcPubSubClient) Close() error {
	return s.client.Close()
}

func (s *gcPubSubClient) getSubscription(dest *Resource) (*pubsub.Subscription, error) {
	if dest.Name == "" && dest.URL == "" {
		return nil, fmt.Errorf("subscription name and URL was empty, expected /projects/[PROJECT ID]/subscriptions/[SUBSCRIPTION] URL or subscription name")
	}
	if dest.projectID == "" {
		return s.client.Subscription(dest.Name), nil
	}
	return s.client.SubscriptionInProject(dest.Name, dest.projectID), nil
}

func (s *gcPubSubClient) getTopic(dest *Resource) (*pubsub.Topic, error) {
	if dest.Name == "" && dest.URL == "" {
		return nil, fmt.Errorf("subscription name and URL was empty, expected /projects/[PROJECT ID]/topics/[SUBSCRIPTION] URL or topic name")
	}
	if dest.projectID == "" {
		return s.client.Topic(dest.Name), nil
	}
	return s.client.TopicInProject(dest.Name, dest.projectID), nil
}

func newCloudPubSub(credConfig *cred.Config, URL string, timeout time.Duration) (Client, error) {
	ctx := context.Background()
	jwtConfig, err := credConfig.NewJWTConfig(pubsub.ScopePubSub)
	if err != nil {
		return nil, err
	}
	var opts = []option.ClientOption{
		option.WithTokenSource(jwtConfig.TokenSource(ctx)),
	}
	client, err := pubsub.NewClient(ctx, credConfig.ProjectID, opts...)
	if err != nil {
		return nil, fmt.Errorf("failed to create pubsub client: %v", err)
	}
	var projectID = extractSubPath(URL, "project")
	if projectID == "" {
		projectID = credConfig.ProjectID
	}

	var service = &gcPubSubClient{
		client:    client,
		ctx:       ctx,
		timeout:   timeout,
		projectId: projectID,
	}
	return service, nil
}
<|endoftext|>


<commit_before>package rain

import (
	"encoding/binary"
	"io"
	"net"
	"time"

	"github.com/cenkalti/log"
)

// All current implementations use 2^14 (16 kiB), and close connections which request an amount greater than that.
const blockSize = 16 * 1024

// http://www.bittorrent.org/beps/bep_0020.html
var peerIDPrefix = []byte("-RN0001-")

type peerID [20]byte

const bitTorrent10pstrLen = 19

var bitTorrent10pstr = []byte("BitTorrent protocol")

func (r *Rain) servePeerConn(conn net.Conn) {
	defer conn.Close()
	log.Debugln("Serving peer", conn.RemoteAddr())

	// Give a minute for completing handshake.
	err := conn.SetDeadline(time.Now().Add(time.Minute))
	if err != nil {
		return
	}

	var d *download

	resultC := make(chan interface{}, 2)
	go r.readHandShake(conn, resultC)

	// Send handshake as soon as you see info_hash.
	i := <-resultC
	switch res := i.(type) {
	case *infoHash:
		// Do not continue if we don't have a torrent with this infoHash.
		r.downloadsM.Lock()
		var ok bool
		if d, ok = r.downloads[*res]; !ok {
			log.Error("unexpected info_hash")
			r.downloadsM.Unlock()
			return
		}
		r.downloadsM.Unlock()

		if err = r.sendHandShake(conn, res); err != nil {
			log.Error(err)
			return
		}
	case error:
		log.Error(res)
		return
	}

	i = <-resultC
	switch res := i.(type) {
	case *peerID:
		if *res == *r.peerID {
			log.Warning("Rejected own connection: server")
			return
		}
		// TODO save peer_id
	case error:
		log.Error(res)
		return
	}

	log.Debugln("servePeerConn: Handshake completed", conn.RemoteAddr())
	p := peerConn{conn: conn, dl: d}
	p.readLoop()
}

func (r *Rain) connectToPeer(p *Peer, d *download) {
	log.Debugln("Connecting to peer", p.TCPAddr())

	conn, err := net.DialTCP("tcp4", nil, p.TCPAddr())
	if err != nil {
		log.Error(err)
		return
	}
	defer conn.Close()

	log.Infoln("Connected to peer", conn.RemoteAddr())

	// Give a minute for completing handshake.
	err = conn.SetDeadline(time.Now().Add(time.Minute))
	if err != nil {
		return
	}

	err = r.sendHandShake(conn, &d.TorrentFile.InfoHash)
	if err != nil {
		log.Error(err)
		return
	}

	ih, id, err := r.readHandShakeBlocking(conn)
	if err != nil {
		log.Error(err)
		return
	}
	if *ih != d.TorrentFile.InfoHash {
		log.Error("unexpected info_hash")
		return
	}
	if *id == *r.peerID {
		log.Warning("Rejected own connection: client")
		return
	}

	log.Debugln("connectToPeer: Handshake completed", conn.RemoteAddr())
	pc := peerConn{conn: conn, dl: d}
	pc.readLoop()
}

const (
	choke = iota
	unchoke
	interested
	notInterested
	have
	bitfield
	request
	piece
	cancel
)

type peerConn struct {
	conn     net.Conn
	dl       *download
	bitfield *BitField
}

// readLoop processes incoming messages after handshake.
func (p *peerConn) readLoop() {
	log.Debugln("Communicating peer", p.conn.RemoteAddr())

	first := true
	buf := make([]byte, blockSize)
	for {
		err := p.conn.SetDeadline(time.Now().Add(3 * time.Minute))
		if err != nil {
			log.Error(err)
			return
		}

		var length uint32
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			log.Error(err)
			return
		}

		if length == 0 { // keepAlive
			log.Debug("received keep-alive")
			err = p.conn.SetDeadline(time.Now().Add(3 * time.Minute))
			if err != nil {
				log.Error(err)
				return
			}
			continue
		}

		var msgType byte
		err = binary.Read(p.conn, binary.BigEndian, &msgType)
		if err != nil {
			log.Error(err)
			return
		}
		length--
		log.Debug("Received message of type ", msgType)

		switch msgType {
		case choke:
		case unchoke:
		case interested:
		case notInterested:
		case have:
			var i int32
			err = binary.Read(p.conn, binary.BigEndian, &i)
			if err != nil {
				log.Error(err)
				return
			}
			p.bitfield.Set(int64(i))
			log.Debugln("Peer", p.conn.RemoteAddr(), "has piece", i)
			log.Debugln("new bitfield:", p.bitfield.Hex())
		case bitfield:
			if !first {
				log.Error("bitfield can only be sent after handshake")
				return
			}

			_, err = io.LimitReader(p.conn, int64(length)).Read(buf)
			if err != nil {
				log.Error(err)
				return
			}

			div, mod := divMod(p.dl.TorrentFile.TotalLength, p.dl.TorrentFile.Info.PieceLength)
			if mod != 0 {
				div++
			}
			log.Debugln("Torrent contains", div, "pieces")
			b := NewBitField(buf, div)
			log.Debugln("Received bitfield:", b.Hex())
			p.bitfield = &b
		case request:
		case piece:
		case cancel:
		}

		first = false
	}
}
<commit_msg>handle choking and interested messages<commit_after>package rain

import (
	"encoding/binary"
	"io"
	"net"
	"time"

	"github.com/cenkalti/log"
)

// All current implementations use 2^14 (16 kiB), and close connections which request an amount greater than that.
const blockSize = 16 * 1024

// http://www.bittorrent.org/beps/bep_0020.html
var peerIDPrefix = []byte("-RN0001-")

type peerID [20]byte

const bitTorrent10pstrLen = 19

var bitTorrent10pstr = []byte("BitTorrent protocol")

func (r *Rain) servePeerConn(conn net.Conn) {
	defer conn.Close()
	log.Debugln("Serving peer", conn.RemoteAddr())

	// Give a minute for completing handshake.
	err := conn.SetDeadline(time.Now().Add(time.Minute))
	if err != nil {
		return
	}

	var d *download

	resultC := make(chan interface{}, 2)
	go r.readHandShake(conn, resultC)

	// Send handshake as soon as you see info_hash.
	i := <-resultC
	switch res := i.(type) {
	case *infoHash:
		// Do not continue if we don't have a torrent with this infoHash.
		r.downloadsM.Lock()
		var ok bool
		if d, ok = r.downloads[*res]; !ok {
			log.Error("unexpected info_hash")
			r.downloadsM.Unlock()
			return
		}
		r.downloadsM.Unlock()

		if err = r.sendHandShake(conn, res); err != nil {
			log.Error(err)
			return
		}
	case error:
		log.Error(res)
		return
	}

	i = <-resultC
	switch res := i.(type) {
	case *peerID:
		if *res == *r.peerID {
			log.Debug("Rejected own connection: server")
			return
		}
		// TODO save peer_id
	case error:
		log.Error(res)
		return
	}

	log.Debugln("servePeerConn: Handshake completed", conn.RemoteAddr())
	p := newPeerConn(conn, d)
	p.readLoop()
}

func (r *Rain) connectToPeer(p *Peer, d *download) {
	log.Debugln("Connecting to peer", p.TCPAddr())

	conn, err := net.DialTCP("tcp4", nil, p.TCPAddr())
	if err != nil {
		log.Error(err)
		return
	}
	defer conn.Close()

	log.Infoln("Connected to peer", conn.RemoteAddr())

	// Give a minute for completing handshake.
	err = conn.SetDeadline(time.Now().Add(time.Minute))
	if err != nil {
		return
	}

	err = r.sendHandShake(conn, &d.TorrentFile.InfoHash)
	if err != nil {
		log.Error(err)
		return
	}

	ih, id, err := r.readHandShakeBlocking(conn)
	if err != nil {
		log.Error(err)
		return
	}
	if *ih != d.TorrentFile.InfoHash {
		log.Error("unexpected info_hash")
		return
	}
	if *id == *r.peerID {
		log.Debug("Rejected own connection: client")
		return
	}

	log.Debugln("connectToPeer: Handshake completed", conn.RemoteAddr())
	pc := newPeerConn(conn, d)
	pc.readLoop()
}

const (
	choke = iota
	unchoke
	interested
	notInterested
	have
	bitfield
	request
	piece
	cancel
)

type peerConn struct {
	conn           net.Conn
	dl             *download
	bitfield       BitField             // on remote
	amChoking      bool                 // this client is choking the peer
	amInterested   bool                 // this client is interested in the peer
	peerChoking    bool                 // peer is choking this client
	peerInterested bool                 // peer is interested in this client
	peerRequests   map[uint64]bool      // What remote peer requested
	ourRequests    map[uint64]time.Time // What we requested, when we requested it
}

func newPeerConn(conn net.Conn, d *download) *peerConn {
	div, mod := divMod(d.TorrentFile.TotalLength, d.TorrentFile.Info.PieceLength)
	if mod != 0 {
		div++
	}
	log.Debugln("Torrent contains", div, "pieces")
	return &peerConn{
		conn:        conn,
		dl:          d,
		bitfield:    NewBitField(nil, div),
		amChoking:   true,
		peerChoking: true,
	}
}

const connReadTimeout = 3 * time.Minute

// readLoop processes incoming messages after handshake.
func (p *peerConn) readLoop() {
	log.Debugln("Communicating peer", p.conn.RemoteAddr())

	first := true
	buf := make([]byte, blockSize)
	for {
		err := p.conn.SetReadDeadline(time.Now().Add(connReadTimeout))
		if err != nil {
			log.Error(err)
			return
		}

		var length uint32
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			log.Error(err)
			return
		}

		if length == 0 { // keepAlive
			log.Debug("received keep-alive")
			continue
		}

		var msgType byte
		err = binary.Read(p.conn, binary.BigEndian, &msgType)
		if err != nil {
			log.Error(err)
			return
		}
		length--
		log.Debug("Received message of type ", msgType)

		switch msgType {
		case choke:
			p.peerChoking = true
		case unchoke:
			p.peerChoking = false
		case interested:
			p.peerInterested = true
		case notInterested:
			p.peerInterested = false
		case have:
			var i int32
			err = binary.Read(p.conn, binary.BigEndian, &i)
			if err != nil {
				log.Error(err)
				return
			}
			p.bitfield.Set(int64(i))
			log.Debugln("Peer", p.conn.RemoteAddr(), "has piece", i)
			log.Debugln("new bitfield:", p.bitfield.Hex())
		case bitfield:
			if !first {
				log.Error("bitfield can only be sent after handshake")
				return
			}

			if int64(length) != int64(len(p.bitfield.Bytes())) {
				log.Error("invalid bitfield length")
				return
			}

			_, err = io.LimitReader(p.conn, int64(length)).Read(buf)
			if err != nil {
				log.Error(err)
				return
			}

			p.bitfield = NewBitField(buf, p.bitfield.Len())
			log.Debugln("Received bitfield:", p.bitfield.Hex())
		case request:
		case piece:
		case cancel:
		}

		first = false
	}
}
<|endoftext|>