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

package shell32

import (
    "syscall"
    "unsafe"
    . "w32"
)

var (
    lib uintptr

    procSHBrowseForFolder   uintptr
    procSHGetPathFromIDList uintptr
)

func init() {
    lib = LoadLib("shell32.dll")

    procSHBrowseForFolder = GetProcAddr(lib, "SHBrowseForFolderW")
    procSHGetPathFromIDList = GetProcAddr(lib, "SHGetPathFromIDListW")
}

func SHBrowseForFolder(bi *BROWSEINFO) uintptr {
    ret, _, _ := syscall.Syscall(procSHBrowseForFolder, 1,
        uintptr(unsafe.Pointer(bi)),
        0,
        0)

    return ret
}

func SHGetPathFromIDList(idl uintptr) string {
    buf := make([]uint16, 1024)
    syscall.Syscall(procSHGetPathFromIDList, 2,
        idl,
        uintptr(unsafe.Pointer(&buf[0])),
        0)

    return syscall.UTF16ToString(buf)
}
<commit_msg>Add DragAcceptFiles,DragQueryFile,DragQueryPoint ,DragFinish.<commit_after>// Copyright 2010 The W32 Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package shell32

import (
    "syscall"
    "unsafe"
    . "w32"
)

var (
    lib uintptr

    procSHBrowseForFolder   uintptr
    procSHGetPathFromIDList uintptr
    procDragAcceptFiles     uintptr
    procDragQueryFile       uintptr
    procDragQueryPoint      uintptr
    procDragFinish          uintptr
)

func init() {
    lib = LoadLib("shell32.dll")

    procSHBrowseForFolder = GetProcAddr(lib, "SHBrowseForFolderW")
    procSHGetPathFromIDList = GetProcAddr(lib, "SHGetPathFromIDListW")
    procDragAcceptFiles = GetProcAddr(lib, "DragAcceptFiles")
    procDragQueryFile = GetProcAddr(lib, "DragQueryFileW")
    procDragQueryPoint = GetProcAddr(lib, "DragQueryPoint")
    procDragFinish = GetProcAddr(lib, "DragFinish")
}

func SHBrowseForFolder(bi *BROWSEINFO) uintptr {
    ret, _, _ := syscall.Syscall(procSHBrowseForFolder, 1,
        uintptr(unsafe.Pointer(bi)),
        0,
        0)

    return ret
}

func SHGetPathFromIDList(idl uintptr) string {
    buf := make([]uint16, 1024)
    syscall.Syscall(procSHGetPathFromIDList, 2,
        idl,
        uintptr(unsafe.Pointer(&buf[0])),
        0)

    return syscall.UTF16ToString(buf)
}

func DragAcceptFiles(hwnd HWND, accept bool) {
    syscall.Syscall(procDragAcceptFiles, 2,
        uintptr(hwnd),
        uintptr(BoolToBOOL(accept)),
        0)
}

func DragQueryFile(hDrop HDROP, iFile uint) (fileName string, fileCount uint) {
    ret, _, _ := syscall.Syscall6(procDragQueryFile, 4,
        uintptr(hDrop),
        uintptr(iFile),
        0,
        0,
        0,
        0)

    fileCount = uint(ret)

    if iFile != 0xFFFFFFFF {
        buf := make([]uint16, fileCount+1)

        ret, _, _ := syscall.Syscall6(procDragQueryFile, 4,
            uintptr(hDrop),
            uintptr(iFile),
            uintptr(unsafe.Pointer(&buf[0])),
            uintptr(fileCount+1),
            0,
            0)

        if ret == 0 {
            panic("Invoke DragQueryFile error.")
        }

        fileName = syscall.UTF16ToString(buf)
    }

    return
}

func DragQueryPoint(hDrop HDROP) (x, y int, isClientArea bool) {
    var pt POINT
    ret, _, _ := syscall.Syscall(procDragQueryPoint, 2,
        uintptr(hDrop),
        uintptr(unsafe.Pointer(&pt)),
        0)

    isClientArea = false
    if ret == 1 {
        isClientArea = true
    }

    x = pt.X
    y = pt.Y

    return
}

func DragFinish(hDrop HDROP) {
    syscall.Syscall(procDragFinish, 1,
        uintptr(hDrop),
        0,
        0)
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Piprate Limited
//
// 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 ld

import (
	"sort"
)

// UsagesNode is a helper class for node usages
type UsagesNode struct {
	node     *NodeMapNode
	property string
	value    map[string]interface{}
}

// NewUsagesNode creates a new instance of UsagesNode
func NewUsagesNode(node *NodeMapNode, property string, value map[string]interface{}) *UsagesNode {
	return &UsagesNode{
		node:     node,
		property: property,
		value:    value,
	}
}

// NodeMapNode
type NodeMapNode struct {
	Values map[string]interface{}
	usages []*UsagesNode
}

// NewNodeMapNode creates a new instance of NodeMapNode.
func NewNodeMapNode(id string) *NodeMapNode {
	return &NodeMapNode{
		Values: map[string]interface{}{"@id": id},
		usages: make([]*UsagesNode, 0),
	}
}

// IsReferencedOnce helps to solve https://github.com/json-ld/json-ld.org/issues/357
// by identifying nodes with just one reference.
func IsReferencedOnce(node *NodeMapNode, referencedOnce map[string]*UsagesNode) bool {
	referencedOnceUsage, present := referencedOnce[node.Values["@id"].(string)]
	return present && referencedOnceUsage != nil
}

// IsWellFormedListNode is a helper function for 4.3.3
func (nmn *NodeMapNode) IsWellFormedListNode() bool {
	keys := 0
	v, containsRdfFirst := nmn.Values[RDFFirst]
	if containsRdfFirst {
		keys++
		vList, isList := v.([]interface{})
		if !(isList && len(vList) == 1) {
			return false
		}
	}
	v, containsRdfRest := nmn.Values[RDFRest]
	if containsRdfRest {
		keys++
		vList, isList := v.([]interface{})
		if !(isList && len(vList) == 1) {
			return false
		}
	}
	v, containsType := nmn.Values["@type"]
	if containsType {
		keys++
		vList, isList := v.([]interface{})
		if !(isList && len(vList) == 1 && vList[0] == RDFList) {
			return false
		}
	}
	// TODO: SPEC: 4.3.3 has no mention of @id
	_, containsID := nmn.Values["@id"]
	if containsID {
		keys++
	}
	if keys < len(nmn.Values) {
		return false
	}
	return true
}

// Serialize returns this node without the usages variable
func (nmn *NodeMapNode) Serialize() map[string]interface{} {
	rval := make(map[string]interface{}, len(nmn.Values))
	for k, v := range nmn.Values {
		rval[k] = v
	}
	return rval
}

// FromRDF converts RDF statements into JSON-LD.
// Returns a list of JSON-LD objects found in the given dataset.
func (api *JsonLdApi) FromRDF(dataset *RDFDataset, opts *JsonLdOptions) ([]interface{}, error) {
	// 1)
	defaultGraph := make(map[string]*NodeMapNode)
	// 2)
	graphMap := make(map[string]map[string]*NodeMapNode)
	graphMap["@default"] = defaultGraph
	referencedOnceMap := make(map[string]*UsagesNode)

	// 3/3.1)
	for name, graph := range dataset.Graphs {
		// 3.2+3.4)
		nodeMap, present := graphMap[name]
		if !present {
			nodeMap = make(map[string]*NodeMapNode)
			graphMap[name] = nodeMap
		}

		// 3.3)
		if _, present := defaultGraph[name]; name != "@default" && !present {
			defaultGraph[name] = NewNodeMapNode(name)
		}

		// 3.5)
		for _, triple := range graph {
			subject := triple.Subject.GetValue()
			predicate := triple.Predicate.GetValue()
			object := triple.Object

			// 3.5.1+3.5.2)
			node, present := nodeMap[subject]
			if !present {
				node = NewNodeMapNode(subject)
				nodeMap[subject] = node
			}

			// 3.5.3)
			_, containsObject := nodeMap[object.GetValue()]
			if (IsIRI(object) || IsBlankNode(object)) && !containsObject {
				nodeMap[object.GetValue()] = NewNodeMapNode(object.GetValue())
			}

			// 3.5.4)
			if predicate == RDFType && (IsIRI(object) || IsBlankNode(object)) && !opts.UseRdfType {
				MergeValue(node.Values, "@type", object.GetValue())
				continue
			}

			// 3.5.5)
			value, _ := rdfToObject(object, opts.UseNativeTypes)

			// 3.5.6+7)
			MergeValue(node.Values, predicate, value)

			// 3.5.8)
			if IsBlankNode(object) || IsIRI(object) {
				// track rdf:nil uniquely per graph
				if object.GetValue() == RDFNil {
					// 3.5.8.1-3)
					n := nodeMap[object.GetValue()]
					n.usages = append(n.usages, NewUsagesNode(node, predicate, value))
				} else if _, present := referencedOnceMap[object.GetValue()]; present {
					referencedOnceMap[object.GetValue()] = nil
				} else {
					// track single reference
					referencedOnceMap[object.GetValue()] = NewUsagesNode(node, predicate, value)
				}
			}
		}
	}

	// 4)
	for _, graph := range graphMap {
		// 4.1), 4.2)
		nilNode, present := graph[RDFNil]
		if !present {
			continue
		}
		// 4.3)
		for _, usage := range nilNode.usages {
			// 4.3.1)
			node := usage.node
			property := usage.property
			head := usage.value
			// 4.3.2)
			list := make([]interface{}, 0)
			listNodes := make([]string, 0)
			// 4.3.3)
			for property == RDFRest && IsReferencedOnce(node, referencedOnceMap) && node.IsWellFormedListNode() {
				// 4.3.3.1)
				list = append(list, node.Values[RDFFirst].([]interface{})[0])
				// 4.3.3.2)
				listNodes = append(listNodes, node.Values["@id"].(string))
				// 4.3.3.3)
				nodeUsage := referencedOnceMap[node.Values["@id"].(string)]
				// 4.3.3.4)
				node = nodeUsage.node
				property = nodeUsage.property
				head = nodeUsage.value
				// 4.3.3.5)
				if !IsBlankNodeValue(node.Values) {
					break
				}
			}

			// 4.3.4)
			if property == RDFFirst {
				// 4.3.4.1)
				if node.Values["@id"] == RDFNil {
					continue
				}
				// 4.3.4.3)
				headID := head["@id"].(string)
				// 4.3.4.4-5)
				head = graph[headID].Values[RDFRest].([]interface{})[0].(map[string]interface{})
				// 4.3.4.6)
				list = list[0 : len(list)-1]
				listNodes = listNodes[0 : len(listNodes)-1]
			}
			// 4.3.5)
			delete(head, "@id")
			// 4.3.6)
			// reverse the list
			for i, j := 0, len(list)-1; i < j; i, j = i+1, j-1 {
				list[i], list[j] = list[j], list[i]
			}
			// 4.3.7)
			head["@list"] = list
			// 4.3.8)
			for _, nodeID := range listNodes {
				delete(graph, nodeID)
			}
		}
	}

	// 5)
	result := make([]interface{}, 0)

	// 6)
	ids := make([]string, 0)
	for k := range defaultGraph {
		ids = append(ids, k)
	}
	sort.Strings(ids)
	for _, subject := range ids {
		node := defaultGraph[subject]
		// 6.1)
		subjectMap, containsSubj := graphMap[subject]
		if containsSubj {
			// 6.1.1)
			graph := make([]interface{}, 0)
			// 6.1.2)
			keys := make([]string, 0)
			for k := range subjectMap {
				keys = append(keys, k)
			}
			sort.Strings(keys)

			for _, s := range keys {
				n := subjectMap[s]
				_, containsID := n.Values["@id"]
				if len(n.Values) == 1 && containsID {
					continue
				}
				graph = append(graph, n.Serialize())
			}
			node.Values["@graph"] = graph
		}
		// 6.2)
		_, containsID := node.Values["@id"]
		if len(node.Values) == 1 && containsID {
			continue
		}
		result = append(result, node.Serialize())
	}

	return result, nil
}
<commit_msg>Support list of lists in FromRdf algorithm<commit_after>// Copyright 2015-2017 Piprate Limited
//
// 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 ld

import (
	"sort"
)

// UsagesNode is a helper class for node usages
type UsagesNode struct {
	node     *NodeMapNode
	property string
	value    map[string]interface{}
}

// NewUsagesNode creates a new instance of UsagesNode
func NewUsagesNode(node *NodeMapNode, property string, value map[string]interface{}) *UsagesNode {
	return &UsagesNode{
		node:     node,
		property: property,
		value:    value,
	}
}

// NodeMapNode
type NodeMapNode struct {
	Values map[string]interface{}
	usages []*UsagesNode
}

// NewNodeMapNode creates a new instance of NodeMapNode.
func NewNodeMapNode(id string) *NodeMapNode {
	return &NodeMapNode{
		Values: map[string]interface{}{"@id": id},
		usages: make([]*UsagesNode, 0),
	}
}

// IsReferencedOnce helps to solve https://github.com/json-ld/json-ld.org/issues/357
// by identifying nodes with just one reference.
func IsReferencedOnce(node *NodeMapNode, referencedOnce map[string]*UsagesNode) bool {
	referencedOnceUsage, present := referencedOnce[node.Values["@id"].(string)]
	return present && referencedOnceUsage != nil
}

// IsWellFormedListNode is a helper function for 4.3.3
func (nmn *NodeMapNode) IsWellFormedListNode() bool {
	keys := 0
	v, containsRdfFirst := nmn.Values[RDFFirst]
	if containsRdfFirst {
		keys++
		vList, isList := v.([]interface{})
		if !(isList && len(vList) == 1) {
			return false
		}
	}
	v, containsRdfRest := nmn.Values[RDFRest]
	if containsRdfRest {
		keys++
		vList, isList := v.([]interface{})
		if !(isList && len(vList) == 1) {
			return false
		}
	}
	v, containsType := nmn.Values["@type"]
	if containsType {
		keys++
		vList, isList := v.([]interface{})
		if !(isList && len(vList) == 1 && vList[0] == RDFList) {
			return false
		}
	}
	// TODO: SPEC: 4.3.3 has no mention of @id
	_, containsID := nmn.Values["@id"]
	if containsID {
		keys++
	}
	if keys < len(nmn.Values) {
		return false
	}
	return true
}

// Serialize returns this node without the usages variable
func (nmn *NodeMapNode) Serialize() map[string]interface{} {
	rval := make(map[string]interface{}, len(nmn.Values))
	for k, v := range nmn.Values {
		rval[k] = v
	}
	return rval
}

// FromRDF converts RDF statements into JSON-LD.
// Returns a list of JSON-LD objects found in the given dataset.
func (api *JsonLdApi) FromRDF(dataset *RDFDataset, opts *JsonLdOptions) ([]interface{}, error) {
	// 1)
	defaultGraph := make(map[string]*NodeMapNode)
	// 2)
	graphMap := make(map[string]map[string]*NodeMapNode)
	graphMap["@default"] = defaultGraph
	referencedOnceMap := make(map[string]*UsagesNode)

	// 3/3.1)
	for name, graph := range dataset.Graphs {
		// 3.2+3.4)
		nodeMap, present := graphMap[name]
		if !present {
			nodeMap = make(map[string]*NodeMapNode)
			graphMap[name] = nodeMap
		}

		// 3.3)
		if _, present := defaultGraph[name]; name != "@default" && !present {
			defaultGraph[name] = NewNodeMapNode(name)
		}

		// 3.5)
		for _, triple := range graph {
			subject := triple.Subject.GetValue()
			predicate := triple.Predicate.GetValue()
			object := triple.Object

			// 3.5.1+3.5.2)
			node, present := nodeMap[subject]
			if !present {
				node = NewNodeMapNode(subject)
				nodeMap[subject] = node
			}

			// 3.5.3)
			_, containsObject := nodeMap[object.GetValue()]
			if (IsIRI(object) || IsBlankNode(object)) && !containsObject {
				nodeMap[object.GetValue()] = NewNodeMapNode(object.GetValue())
			}

			// 3.5.4)
			if predicate == RDFType && (IsIRI(object) || IsBlankNode(object)) && !opts.UseRdfType {
				MergeValue(node.Values, "@type", object.GetValue())
				continue
			}

			// 3.5.5)
			value, err := rdfToObject(object, opts.UseNativeTypes)
			if err != nil {
				return nil, err
			}

			// 3.5.6+7)
			MergeValue(node.Values, predicate, value)

			// 3.5.8)
			if IsBlankNode(object) || IsIRI(object) {
				// track rdf:nil uniquely per graph
				if object.GetValue() == RDFNil {
					// 3.5.8.1-3)
					n := nodeMap[object.GetValue()]
					n.usages = append(n.usages, NewUsagesNode(node, predicate, value))
				} else if _, present := referencedOnceMap[object.GetValue()]; present {
					referencedOnceMap[object.GetValue()] = nil
				} else {
					// track single reference
					referencedOnceMap[object.GetValue()] = NewUsagesNode(node, predicate, value)
				}
			}
		}
	}

	// 4)
	for _, graph := range graphMap {
		// 4.1), 4.2)
		nilNode, present := graph[RDFNil]
		if !present {
			continue
		}
		// 4.3)
		for _, usage := range nilNode.usages {
			// 4.3.1)
			node := usage.node
			property := usage.property
			head := usage.value
			// 4.3.2)
			list := make([]interface{}, 0)
			listNodes := make([]string, 0)
			// 4.3.3)
			for property == RDFRest && IsReferencedOnce(node, referencedOnceMap) && node.IsWellFormedListNode() {
				// 4.3.3.1)
				list = append(list, node.Values[RDFFirst].([]interface{})[0])
				// 4.3.3.2)
				listNodes = append(listNodes, node.Values["@id"].(string))
				// 4.3.3.3)
				nodeUsage := referencedOnceMap[node.Values["@id"].(string)]
				// 4.3.3.4)
				node = nodeUsage.node
				property = nodeUsage.property
				head = nodeUsage.value
				// if node is not a blank node, then list head found
				if !IsBlankNodeValue(node.Values) {
					break
				}
			}

			// 4.3.5)
			delete(head, "@id")
			// 4.3.6)
			// reverse the list
			for i, j := 0, len(list)-1; i < j; i, j = i+1, j-1 {
				list[i], list[j] = list[j], list[i]
			}
			// 4.3.7)
			head["@list"] = list
			// 4.3.8)
			for _, nodeID := range listNodes {
				delete(graph, nodeID)
			}
		}
	}

	// 5)
	result := make([]interface{}, 0)

	// 6)
	ids := make([]string, 0)
	for k := range defaultGraph {
		ids = append(ids, k)
	}
	sort.Strings(ids)
	for _, subject := range ids {
		node := defaultGraph[subject]
		// 6.1)
		subjectMap, containsSubj := graphMap[subject]
		if containsSubj {
			// 6.1.1)
			graph := make([]interface{}, 0)
			// 6.1.2)
			keys := make([]string, 0)
			for k := range subjectMap {
				keys = append(keys, k)
			}
			sort.Strings(keys)

			for _, s := range keys {
				n := subjectMap[s]
				_, containsID := n.Values["@id"]
				if len(n.Values) == 1 && containsID {
					continue
				}
				graph = append(graph, n.Serialize())
			}
			node.Values["@graph"] = graph
		}
		// 6.2)
		_, containsID := node.Values["@id"]
		if len(node.Values) == 1 && containsID {
			continue
		}
		result = append(result, node.Serialize())
	}

	return result, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes 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.
*/

package cloudresource

import (
	"context"
	"fmt"
	"sync"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/wait"
	cloudprovider "k8s.io/cloud-provider"

	"k8s.io/klog/v2"
)

// SyncManager is an interface for making requests to a cloud provider
type SyncManager interface {
	Run(stopCh <-chan struct{})
	NodeAddresses() ([]v1.NodeAddress, error)
}

var _ SyncManager = &cloudResourceSyncManager{}

type cloudResourceSyncManager struct {
	// Cloud provider interface.
	cloud cloudprovider.Interface
	// Sync period
	syncPeriod time.Duration

	nodeAddressesMonitor *sync.Cond
	nodeAddressesErr     error
	nodeAddresses        []v1.NodeAddress

	nodeName types.NodeName
}

// NewSyncManager creates a manager responsible for collecting resources from a
// cloud provider through requests that are sensitive to timeouts and hanging
func NewSyncManager(cloud cloudprovider.Interface, nodeName types.NodeName, syncPeriod time.Duration) SyncManager {
	return &cloudResourceSyncManager{
		cloud:      cloud,
		syncPeriod: syncPeriod,
		nodeName:   nodeName,
		// nodeAddressesMonitor is a monitor that guards a result (nodeAddresses,
		// nodeAddressesErr) of the sync loop under the condition that a result has
		// been saved at least once. The semantics here are:
		//
		// * Readers of the result will wait on the monitor until the first result
		//   has been saved.
		// * The sync loop (i.e. the only writer), will signal all waiters every
		//   time it updates the result.
		nodeAddressesMonitor: sync.NewCond(&sync.Mutex{}),
	}
}

// NodeAddresses waits for the first sync loop to run. If no successful syncs
// have run, it will return the most recent error. If node addresses have been
// synced successfully, it will return the list of node addresses from the most
// recent successful sync.
func (m *cloudResourceSyncManager) NodeAddresses() ([]v1.NodeAddress, error) {
	m.nodeAddressesMonitor.L.Lock()
	defer m.nodeAddressesMonitor.L.Unlock()
	// wait until there is something
	for {
		if addrs, err := m.nodeAddresses, m.nodeAddressesErr; len(addrs) > 0 || err != nil {
			return addrs, err
		}
		klog.V(5).InfoS("Waiting for cloud provider to provide node addresses")
		m.nodeAddressesMonitor.Wait()
	}
}

// getNodeAddresses calls the cloud provider to get a current list of node addresses.
func (m *cloudResourceSyncManager) getNodeAddresses() ([]v1.NodeAddress, error) {
	// TODO(roberthbailey): Can we do this without having credentials to talk to
	// the cloud provider?
	// TODO(justinsb): We can if CurrentNodeName() was actually CurrentNode() and
	// returned an interface.
	// TODO: If IP addresses couldn't be fetched from the cloud provider, should
	// kubelet fallback on the other methods for getting the IP below?
	instances, ok := m.cloud.Instances()
	if !ok {
		return nil, fmt.Errorf("failed to get instances from cloud provider")
	}
	return instances.NodeAddresses(context.TODO(), m.nodeName)
}

func (m *cloudResourceSyncManager) syncNodeAddresses() {
	klog.V(5).InfoS("Requesting node addresses from cloud provider for node","nodeName", m.nodeName)

	addrs, err := m.getNodeAddresses()

	m.nodeAddressesMonitor.L.Lock()
	defer m.nodeAddressesMonitor.L.Unlock()
	defer m.nodeAddressesMonitor.Broadcast()

	if err != nil {
		klog.V(2).InfoS("Node addresses from cloud provider for node not collected","nodeName", m.nodeName, "err", err)

		if len(m.nodeAddresses) > 0 {
			// in the event that a sync loop fails when a previous sync had
			// succeeded, continue to use the old addresses.
			return
		}

		m.nodeAddressesErr = fmt.Errorf("failed to get node address from cloud provider: %v", err)
		return
	}

	klog.V(5).InfoS("Node addresses from cloud provider for node collected","nodeName", m.nodeName)
	m.nodeAddressesErr = nil
	m.nodeAddresses = addrs
}

// Run starts the cloud resource sync manager's sync loop.
func (m *cloudResourceSyncManager) Run(stopCh <-chan struct{}) {
	wait.Until(m.syncNodeAddresses, m.syncPeriod, stopCh)
}
<commit_msg>Apply suggestions from code review<commit_after>/*
Copyright 2018 The Kubernetes 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.
*/

package cloudresource

import (
	"context"
	"fmt"
	"sync"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/wait"
	cloudprovider "k8s.io/cloud-provider"

	"k8s.io/klog/v2"
)

// SyncManager is an interface for making requests to a cloud provider
type SyncManager interface {
	Run(stopCh <-chan struct{})
	NodeAddresses() ([]v1.NodeAddress, error)
}

var _ SyncManager = &cloudResourceSyncManager{}

type cloudResourceSyncManager struct {
	// Cloud provider interface.
	cloud cloudprovider.Interface
	// Sync period
	syncPeriod time.Duration

	nodeAddressesMonitor *sync.Cond
	nodeAddressesErr     error
	nodeAddresses        []v1.NodeAddress

	nodeName types.NodeName
}

// NewSyncManager creates a manager responsible for collecting resources from a
// cloud provider through requests that are sensitive to timeouts and hanging
func NewSyncManager(cloud cloudprovider.Interface, nodeName types.NodeName, syncPeriod time.Duration) SyncManager {
	return &cloudResourceSyncManager{
		cloud:      cloud,
		syncPeriod: syncPeriod,
		nodeName:   nodeName,
		// nodeAddressesMonitor is a monitor that guards a result (nodeAddresses,
		// nodeAddressesErr) of the sync loop under the condition that a result has
		// been saved at least once. The semantics here are:
		//
		// * Readers of the result will wait on the monitor until the first result
		//   has been saved.
		// * The sync loop (i.e. the only writer), will signal all waiters every
		//   time it updates the result.
		nodeAddressesMonitor: sync.NewCond(&sync.Mutex{}),
	}
}

// NodeAddresses waits for the first sync loop to run. If no successful syncs
// have run, it will return the most recent error. If node addresses have been
// synced successfully, it will return the list of node addresses from the most
// recent successful sync.
func (m *cloudResourceSyncManager) NodeAddresses() ([]v1.NodeAddress, error) {
	m.nodeAddressesMonitor.L.Lock()
	defer m.nodeAddressesMonitor.L.Unlock()
	// wait until there is something
	for {
		if addrs, err := m.nodeAddresses, m.nodeAddressesErr; len(addrs) > 0 || err != nil {
			return addrs, err
		}
		klog.V(5).InfoS("Waiting for cloud provider to provide node addresses")
		m.nodeAddressesMonitor.Wait()
	}
}

// getNodeAddresses calls the cloud provider to get a current list of node addresses.
func (m *cloudResourceSyncManager) getNodeAddresses() ([]v1.NodeAddress, error) {
	// TODO(roberthbailey): Can we do this without having credentials to talk to
	// the cloud provider?
	// TODO(justinsb): We can if CurrentNodeName() was actually CurrentNode() and
	// returned an interface.
	// TODO: If IP addresses couldn't be fetched from the cloud provider, should
	// kubelet fallback on the other methods for getting the IP below?
	instances, ok := m.cloud.Instances()
	if !ok {
		return nil, fmt.Errorf("failed to get instances from cloud provider")
	}
	return instances.NodeAddresses(context.TODO(), m.nodeName)
}

func (m *cloudResourceSyncManager) syncNodeAddresses() {
	klog.V(5).InfoS("Requesting node addresses from cloud provider for node", "nodeName", m.nodeName)

	addrs, err := m.getNodeAddresses()

	m.nodeAddressesMonitor.L.Lock()
	defer m.nodeAddressesMonitor.L.Unlock()
	defer m.nodeAddressesMonitor.Broadcast()

	if err != nil {
		klog.V(2).InfoS("Node addresses from cloud provider for node not collected", "nodeName", m.nodeName, "err", err)

		if len(m.nodeAddresses) > 0 {
			// in the event that a sync loop fails when a previous sync had
			// succeeded, continue to use the old addresses.
			return
		}

		m.nodeAddressesErr = fmt.Errorf("failed to get node address from cloud provider: %v", err)
		return
	}

	klog.V(5).InfoS("Node addresses from cloud provider for node collected", "nodeName", m.nodeName)
	m.nodeAddressesErr = nil
	m.nodeAddresses = addrs
}

// Run starts the cloud resource sync manager's sync loop.
func (m *cloudResourceSyncManager) Run(stopCh <-chan struct{}) {
	wait.Until(m.syncNodeAddresses, m.syncPeriod, stopCh)
}
<|endoftext|>


<commit_before>// +build windows

/*
Copyright 2019 The Kubernetes 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.
*/

package dockershim

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"regexp"

	"golang.org/x/sys/windows/registry"

	dockertypes "github.com/docker/docker/api/types"
	dockercontainer "github.com/docker/docker/api/types/container"

	utilfeature "k8s.io/apiserver/pkg/util/feature"
	kubefeatures "k8s.io/kubernetes/pkg/features"
	runtimeapi "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
	"k8s.io/kubernetes/pkg/kubelet/kuberuntime"
)

type containerCreationCleanupInfo struct {
	gMSARegistryValueName string
}

// applyPlatformSpecificDockerConfig applies platform-specific configurations to a dockertypes.ContainerCreateConfig struct.
// The containerCreationCleanupInfo struct it returns will be passed as is to performPlatformSpecificContainerCreationCleanup
// after the container has been created.
func (ds *dockerService) applyPlatformSpecificDockerConfig(request *runtimeapi.CreateContainerRequest, createConfig *dockertypes.ContainerCreateConfig) (*containerCreationCleanupInfo, error) {
	cleanupInfo := &containerCreationCleanupInfo{}

	if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.WindowsGMSA) {
		if err := applyGMSAConfig(request.GetConfig(), createConfig, cleanupInfo); err != nil {
			return nil, err
		}
	}

	return cleanupInfo, nil
}

// applyGMSAConfig looks at the kuberuntime.GMSASpecContainerAnnotationKey container annotation; if present,
// it copies its contents to a unique registry value, and sets a SecurityOpt on the config pointing to that registry value.
// We use registry values instead of files since their location cannot change - as opposed to credential spec files,
// whose location could potentially change down the line, or even be unknown (eg if docker is not installed on the
// C: drive)
// When docker supports passing a credential spec's contents directly, we should switch to using that
// as it will avoid cluttering the registry.
func applyGMSAConfig(config *runtimeapi.ContainerConfig, createConfig *dockertypes.ContainerCreateConfig, cleanupInfo *containerCreationCleanupInfo) error {
	credSpec := config.Annotations[kuberuntime.GMSASpecContainerAnnotationKey]
	if credSpec == "" {
		return nil
	}

	valueName, err := copyGMSACredSpecToRegistryValue(credSpec)
	if err != nil {
		return err
	}

	if createConfig.HostConfig == nil {
		createConfig.HostConfig = &dockercontainer.HostConfig{}
	}

	createConfig.HostConfig.SecurityOpt = append(createConfig.HostConfig.SecurityOpt, "credentialspec=registry://"+valueName)
	cleanupInfo.gMSARegistryValueName = valueName

	return nil
}

const (
	// same as https://github.com/moby/moby/blob/93d994e29c9cc8d81f1b0477e28d705fa7e2cd72/daemon/oci_windows.go#L23
	credentialSpecRegistryLocation = `SOFTWARE\Microsoft\Windows NT\CurrentVersion\Virtualization\Containers\CredentialSpecs`
	// the prefix for the registry values we write GMSA cred specs to
	gMSARegistryValueNamePrefix = "k8s-cred-spec-"
	// the number of random bytes to generate suffixes for registry value names
	gMSARegistryValueNameSuffixRandomBytes = 40
)

// registryKey is an interface wrapper around `registry.Key`,
// listing only the methods we care about here.
// It's mainly useful to easily allow mocking the registry in tests.
type registryKey interface {
	SetStringValue(name, value string) error
	DeleteValue(name string) error
	ReadValueNames(n int) ([]string, error)
	Close() error
}

var registryCreateKeyFunc = func(baseKey registry.Key, path string, access uint32) (registryKey, bool, error) {
	return registry.CreateKey(baseKey, path, access)
}

// randomReader is only meant to ever be overriden for testing purposes,
// same idea as for `registryKey` above
var randomReader = rand.Reader

// gMSARegistryValueNamesRegex is the regex used to detect gMSA cred spec
// registry values in `removeAllGMSARegistryValues` below.
var gMSARegistryValueNamesRegex = regexp.MustCompile(fmt.Sprintf("^%s[0-9a-f]{%d}$", gMSARegistryValueNamePrefix, 2*gMSARegistryValueNameSuffixRandomBytes))

// copyGMSACredSpecToRegistryKey copies the credential specs to a unique registry value, and returns its name.
func copyGMSACredSpecToRegistryValue(credSpec string) (string, error) {
	valueName, err := gMSARegistryValueName()
	if err != nil {
		return "", err
	}

	// write to the registry
	key, _, err := registryCreateKeyFunc(registry.LOCAL_MACHINE, credentialSpecRegistryLocation, registry.SET_VALUE)
	if err != nil {
		return "", fmt.Errorf("unable to open registry key %q: %v", credentialSpecRegistryLocation, err)
	}
	defer key.Close()
	if err = key.SetStringValue(valueName, credSpec); err != nil {
		return "", fmt.Errorf("unable to write into registry value %q/%q: %v", credentialSpecRegistryLocation, valueName, err)
	}

	return valueName, nil
}

// gMSARegistryValueName computes the name of the registry value where to store the GMSA cred spec contents.
// The value's name is a purely random suffix appended to `gMSARegistryValueNamePrefix`.
func gMSARegistryValueName() (string, error) {
	randomSuffix, err := randomString(gMSARegistryValueNameSuffixRandomBytes)

	if err != nil {
		return "", fmt.Errorf("error when generating gMSA registry value name: %v", err)
	}

	return gMSARegistryValueNamePrefix + randomSuffix, nil
}

// randomString returns a random hex string.
func randomString(length int) (string, error) {
	randBytes := make([]byte, length)

	if n, err := randomReader.Read(randBytes); err != nil || n != length {
		if err == nil {
			err = fmt.Errorf("only got %v random bytes, expected %v", n, length)
		}
		return "", fmt.Errorf("unable to generate random string: %v", err)
	}

	return hex.EncodeToString(randBytes), nil
}

// performPlatformSpecificContainerCreationCleanup is responsible for doing any platform-specific cleanup
// after a container creation. Any errors it returns are simply logged, but do not fail the container
// creation.
func (ds *dockerService) performPlatformSpecificContainerCreationCleanup(cleanupInfo *containerCreationCleanupInfo) (errors []error) {
	if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.WindowsGMSA) {
		if err := removeGMSARegistryValue(cleanupInfo); err != nil {
			errors = append(errors, err)
		}
	}

	return
}

func removeGMSARegistryValue(cleanupInfo *containerCreationCleanupInfo) error {
	if cleanupInfo == nil || cleanupInfo.gMSARegistryValueName == "" {
		return nil
	}

	key, _, err := registryCreateKeyFunc(registry.LOCAL_MACHINE, credentialSpecRegistryLocation, registry.SET_VALUE)
	if err != nil {
		return fmt.Errorf("unable to open registry key %q: %v", credentialSpecRegistryLocation, err)
	}
	defer key.Close()
	if err = key.DeleteValue(cleanupInfo.gMSARegistryValueName); err != nil {
		return fmt.Errorf("unable to remove registry value %q/%q: %v", credentialSpecRegistryLocation, cleanupInfo.gMSARegistryValueName, err)
	}

	return nil
}

// platformSpecificContainerCreationInitCleanup is called when dockershim
// is starting, and is meant to clean up any cruft left by previous runs
// creating containers.
// Errors are simply logged, but don't prevent dockershim from starting.
func (ds *dockerService) platformSpecificContainerCreationInitCleanup() (errors []error) {
	if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.WindowsGMSA) {
		errors = removeAllGMSARegistryValues()
	}
	return
}

func removeAllGMSARegistryValues() (errors []error) {
	key, _, err := registryCreateKeyFunc(registry.LOCAL_MACHINE, credentialSpecRegistryLocation, registry.SET_VALUE)
	if err != nil {
		return []error{fmt.Errorf("unable to open registry key %q: %v", credentialSpecRegistryLocation, err)}
	}
	defer key.Close()

	valueNames, err := key.ReadValueNames(0)
	if err != nil {
		return []error{fmt.Errorf("unable to list values under registry key %q: %v", credentialSpecRegistryLocation, err)}
	}

	for _, valueName := range valueNames {
		if gMSARegistryValueNamesRegex.MatchString(valueName) {
			if err = key.DeleteValue(valueName); err != nil {
				errors = append(errors, fmt.Errorf("unable to remove registry value %q/%q: %v", credentialSpecRegistryLocation, valueName, err))
			}
		}
	}

	return
}
<commit_msg>Typo in comment<commit_after>// +build windows

/*
Copyright 2019 The Kubernetes 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.
*/

package dockershim

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"regexp"

	"golang.org/x/sys/windows/registry"

	dockertypes "github.com/docker/docker/api/types"
	dockercontainer "github.com/docker/docker/api/types/container"

	utilfeature "k8s.io/apiserver/pkg/util/feature"
	kubefeatures "k8s.io/kubernetes/pkg/features"
	runtimeapi "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
	"k8s.io/kubernetes/pkg/kubelet/kuberuntime"
)

type containerCreationCleanupInfo struct {
	gMSARegistryValueName string
}

// applyPlatformSpecificDockerConfig applies platform-specific configurations to a dockertypes.ContainerCreateConfig struct.
// The containerCreationCleanupInfo struct it returns will be passed as is to performPlatformSpecificContainerCreationCleanup
// after the container has been created.
func (ds *dockerService) applyPlatformSpecificDockerConfig(request *runtimeapi.CreateContainerRequest, createConfig *dockertypes.ContainerCreateConfig) (*containerCreationCleanupInfo, error) {
	cleanupInfo := &containerCreationCleanupInfo{}

	if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.WindowsGMSA) {
		if err := applyGMSAConfig(request.GetConfig(), createConfig, cleanupInfo); err != nil {
			return nil, err
		}
	}

	return cleanupInfo, nil
}

// applyGMSAConfig looks at the kuberuntime.GMSASpecContainerAnnotationKey container annotation; if present,
// it copies its contents to a unique registry value, and sets a SecurityOpt on the config pointing to that registry value.
// We use registry values instead of files since their location cannot change - as opposed to credential spec files,
// whose location could potentially change down the line, or even be unknown (eg if docker is not installed on the
// C: drive)
// When docker supports passing a credential spec's contents directly, we should switch to using that
// as it will avoid cluttering the registry.
func applyGMSAConfig(config *runtimeapi.ContainerConfig, createConfig *dockertypes.ContainerCreateConfig, cleanupInfo *containerCreationCleanupInfo) error {
	credSpec := config.Annotations[kuberuntime.GMSASpecContainerAnnotationKey]
	if credSpec == "" {
		return nil
	}

	valueName, err := copyGMSACredSpecToRegistryValue(credSpec)
	if err != nil {
		return err
	}

	if createConfig.HostConfig == nil {
		createConfig.HostConfig = &dockercontainer.HostConfig{}
	}

	createConfig.HostConfig.SecurityOpt = append(createConfig.HostConfig.SecurityOpt, "credentialspec=registry://"+valueName)
	cleanupInfo.gMSARegistryValueName = valueName

	return nil
}

const (
	// same as https://github.com/moby/moby/blob/93d994e29c9cc8d81f1b0477e28d705fa7e2cd72/daemon/oci_windows.go#L23
	credentialSpecRegistryLocation = `SOFTWARE\Microsoft\Windows NT\CurrentVersion\Virtualization\Containers\CredentialSpecs`
	// the prefix for the registry values we write GMSA cred specs to
	gMSARegistryValueNamePrefix = "k8s-cred-spec-"
	// the number of random bytes to generate suffixes for registry value names
	gMSARegistryValueNameSuffixRandomBytes = 40
)

// registryKey is an interface wrapper around `registry.Key`,
// listing only the methods we care about here.
// It's mainly useful to easily allow mocking the registry in tests.
type registryKey interface {
	SetStringValue(name, value string) error
	DeleteValue(name string) error
	ReadValueNames(n int) ([]string, error)
	Close() error
}

var registryCreateKeyFunc = func(baseKey registry.Key, path string, access uint32) (registryKey, bool, error) {
	return registry.CreateKey(baseKey, path, access)
}

// randomReader is only meant to ever be overridden for testing purposes,
// same idea as for `registryKey` above
var randomReader = rand.Reader

// gMSARegistryValueNamesRegex is the regex used to detect gMSA cred spec
// registry values in `removeAllGMSARegistryValues` below.
var gMSARegistryValueNamesRegex = regexp.MustCompile(fmt.Sprintf("^%s[0-9a-f]{%d}$", gMSARegistryValueNamePrefix, 2*gMSARegistryValueNameSuffixRandomBytes))

// copyGMSACredSpecToRegistryKey copies the credential specs to a unique registry value, and returns its name.
func copyGMSACredSpecToRegistryValue(credSpec string) (string, error) {
	valueName, err := gMSARegistryValueName()
	if err != nil {
		return "", err
	}

	// write to the registry
	key, _, err := registryCreateKeyFunc(registry.LOCAL_MACHINE, credentialSpecRegistryLocation, registry.SET_VALUE)
	if err != nil {
		return "", fmt.Errorf("unable to open registry key %q: %v", credentialSpecRegistryLocation, err)
	}
	defer key.Close()
	if err = key.SetStringValue(valueName, credSpec); err != nil {
		return "", fmt.Errorf("unable to write into registry value %q/%q: %v", credentialSpecRegistryLocation, valueName, err)
	}

	return valueName, nil
}

// gMSARegistryValueName computes the name of the registry value where to store the GMSA cred spec contents.
// The value's name is a purely random suffix appended to `gMSARegistryValueNamePrefix`.
func gMSARegistryValueName() (string, error) {
	randomSuffix, err := randomString(gMSARegistryValueNameSuffixRandomBytes)

	if err != nil {
		return "", fmt.Errorf("error when generating gMSA registry value name: %v", err)
	}

	return gMSARegistryValueNamePrefix + randomSuffix, nil
}

// randomString returns a random hex string.
func randomString(length int) (string, error) {
	randBytes := make([]byte, length)

	if n, err := randomReader.Read(randBytes); err != nil || n != length {
		if err == nil {
			err = fmt.Errorf("only got %v random bytes, expected %v", n, length)
		}
		return "", fmt.Errorf("unable to generate random string: %v", err)
	}

	return hex.EncodeToString(randBytes), nil
}

// performPlatformSpecificContainerCreationCleanup is responsible for doing any platform-specific cleanup
// after a container creation. Any errors it returns are simply logged, but do not fail the container
// creation.
func (ds *dockerService) performPlatformSpecificContainerCreationCleanup(cleanupInfo *containerCreationCleanupInfo) (errors []error) {
	if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.WindowsGMSA) {
		if err := removeGMSARegistryValue(cleanupInfo); err != nil {
			errors = append(errors, err)
		}
	}

	return
}

func removeGMSARegistryValue(cleanupInfo *containerCreationCleanupInfo) error {
	if cleanupInfo == nil || cleanupInfo.gMSARegistryValueName == "" {
		return nil
	}

	key, _, err := registryCreateKeyFunc(registry.LOCAL_MACHINE, credentialSpecRegistryLocation, registry.SET_VALUE)
	if err != nil {
		return fmt.Errorf("unable to open registry key %q: %v", credentialSpecRegistryLocation, err)
	}
	defer key.Close()
	if err = key.DeleteValue(cleanupInfo.gMSARegistryValueName); err != nil {
		return fmt.Errorf("unable to remove registry value %q/%q: %v", credentialSpecRegistryLocation, cleanupInfo.gMSARegistryValueName, err)
	}

	return nil
}

// platformSpecificContainerCreationInitCleanup is called when dockershim
// is starting, and is meant to clean up any cruft left by previous runs
// creating containers.
// Errors are simply logged, but don't prevent dockershim from starting.
func (ds *dockerService) platformSpecificContainerCreationInitCleanup() (errors []error) {
	if utilfeature.DefaultFeatureGate.Enabled(kubefeatures.WindowsGMSA) {
		errors = removeAllGMSARegistryValues()
	}
	return
}

func removeAllGMSARegistryValues() (errors []error) {
	key, _, err := registryCreateKeyFunc(registry.LOCAL_MACHINE, credentialSpecRegistryLocation, registry.SET_VALUE)
	if err != nil {
		return []error{fmt.Errorf("unable to open registry key %q: %v", credentialSpecRegistryLocation, err)}
	}
	defer key.Close()

	valueNames, err := key.ReadValueNames(0)
	if err != nil {
		return []error{fmt.Errorf("unable to list values under registry key %q: %v", credentialSpecRegistryLocation, err)}
	}

	for _, valueName := range valueNames {
		if gMSARegistryValueNamesRegex.MatchString(valueName) {
			if err = key.DeleteValue(valueName); err != nil {
				errors = append(errors, fmt.Errorf("unable to remove registry value %q/%q: %v", credentialSpecRegistryLocation, valueName, err))
			}
		}
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/nuclio/nuclio/pkg/errors"
	"github.com/nuclio/nuclio/pkg/functionconfig"
	"github.com/nuclio/nuclio/pkg/loggersink"
	nuclioio_client "github.com/nuclio/nuclio/pkg/platform/kube/client/clientset/versioned"
	// load all sinks
	"github.com/nuclio/nuclio/pkg/platformconfig"
	_ "github.com/nuclio/nuclio/pkg/sinks"

	"github.com/nuclio/logger"
	"github.com/v3io/scaler-types"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

// A plugin for github.com/v3io/scaler, allowing to extend it to scale to zero and from zero function resources in k8s
type NuclioResourceScaler struct {
	logger          logger.Logger
	nuclioClientSet nuclioio_client.Interface
	kubeconfigPath  string
	namespace       string
}

// New is called when plugin loaded on scaler, so it's considered "dead code" for the linter
func New(kubeconfigPath string, namespace string) (scaler_types.ResourceScaler, error) { // nolint: deadcode
	platformConfiguration, err := readPlatformConfiguration()
	if err != nil {
		return nil, errors.Wrap(err, "Failed to read platform configuration")
	}

	// create a root logger
	resourceScalerLogger, err := loggersink.CreateSystemLogger("resource-scaler", platformConfiguration)
	if err != nil {
		fmt.Println(err)
		return nil, errors.Wrap(err, "Failed to create logger")
	}

	restConfig, err := getClientConfig(kubeconfigPath)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to get client configuration")
	}

	nuclioClientSet, err := nuclioio_client.NewForConfig(restConfig)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to create nuclio client set")
	}

	resourceScalerLogger.DebugWith("Initialized resource scaler",
		"platformconfig", platformConfiguration,
		"namespace", namespace,
		"kubeconfigPath", kubeconfigPath)

	return &NuclioResourceScaler{
		logger:          resourceScalerLogger,
		nuclioClientSet: nuclioClientSet,
		kubeconfigPath:  kubeconfigPath,
		namespace:       namespace,
	}, nil
}

func (n *NuclioResourceScaler) SetScale(resource scaler_types.Resource, scale int) error {
	if scale == 0 {
		return n.scaleFunctionToZero(n.namespace, string(resource))
	}
	return n.scaleFunctionFromZero(n.namespace, string(resource))
}

func (n *NuclioResourceScaler) GetResources() ([]scaler_types.Resource, error) {
	functions, err := n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(n.namespace).List(metav1.ListOptions{})
	if err != nil {
		return nil, errors.Wrap(err, "Failed to list functions")
	}

	var functionList []scaler_types.Resource

	// build a map of functions with status
	for _, function := range functions.Items {
		functionList = append(functionList, scaler_types.Resource(function.Name))
	}
	return functionList, nil
}

func (n *NuclioResourceScaler) GetConfig() (*scaler_types.ResourceScalerConfig, error) {
	platformConfiguration, err := readPlatformConfiguration()
	if err != nil {
		return nil, errors.Wrap(err, "Failed to read platform configuration")
	}

	scaleInterval, err := time.ParseDuration(platformConfiguration.ScaleToZero.ScalerInterval)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to read scale interval")
	}

	scaleWindow, err := time.ParseDuration(platformConfiguration.ScaleToZero.WindowSize)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to parse scale window")
	}

	pollerInterval, err := time.ParseDuration(platformConfiguration.ScaleToZero.PollerInterval)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to parse poller interval")
	}

	return &scaler_types.ResourceScalerConfig{
		KubeconfigPath: n.kubeconfigPath,
		AutoScalerOptions: scaler_types.AutoScalerOptions{
			Namespace:     n.namespace,
			ScaleInterval: scaleInterval,
			ScaleWindow:   scaleWindow,
			MetricName:    platformConfiguration.ScaleToZero.MetricName,
			Threshold:     0,
		},
		DLXOptions: scaler_types.DLXOptions{
			Namespace:        n.namespace,
			TargetPort:       8080,
			TargetNameHeader: "X-Nuclio-Target",
			TargetPathHeader: "X-Nuclio-Function-Path",
			ListenAddress:    ":8080",
		},
		PollerOptions: scaler_types.PollerOptions{
			MetricInterval: pollerInterval,
			MetricName:     platformConfiguration.ScaleToZero.MetricName,
			Namespace:      n.namespace,
			GroupKind:      "NuclioFunction",
		},
	}, nil
}

func (n *NuclioResourceScaler) scaleFunctionToZero(namespace string, functionName string) error {
	n.logger.DebugWith("Scaling to zero", "functionName", functionName)
	function, err := n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Get(functionName, metav1.GetOptions{})
	if err != nil {
		n.logger.WarnWith("Failed to get nuclio function", "functionName", functionName, "err", err)
		return errors.Wrap(err, "Failed to get function")
	}

	// this has the nice property of disabling hpa as well
	function.Status.State = functionconfig.FunctionStateScaledToZero
	_, err = n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Update(function)
	if err != nil {
		n.logger.WarnWith("Failed to update function", "functionName", functionName, "err", err)
		return errors.Wrap(err, "Failed to update function")
	}
	return nil
}

func (n *NuclioResourceScaler) scaleFunctionFromZero(namespace string, functionName string) error {
	err := n.updateFunctionStatus(namespace, functionName)
	if err != nil {
		return errors.Wrap(err, "Failed to change function status to waitingForResourceConfiguration")
	}
	return n.waitFunctionReadiness(namespace, functionName)
}

func (n *NuclioResourceScaler) updateFunctionStatus(namespace string, functionName string) error {
	function, err := n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Get(functionName, metav1.GetOptions{})
	if err != nil {
		n.logger.WarnWith("Failed to get nuclio function", "functionName", functionName, "err", err)
		return errors.Wrap(err, "Failed to get nuclio function")
	}

	function.Status.State = functionconfig.FunctionStateWaitingForResourceConfiguration
	_, err = n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Update(function)
	if err != nil {
		n.logger.WarnWith("Failed to update function", "functionName", functionName, "err", err)
		return errors.Wrap(err, "Failed to update nuclio function")
	}
	return nil
}

func (n *NuclioResourceScaler) waitFunctionReadiness(namespace string, functionName string) error {
	for {
		function, err := n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Get(functionName, metav1.GetOptions{})
		if err != nil {
			n.logger.WarnWith("Failed to get nuclio function", "functionName", functionName, "err", err)
			return errors.Wrap(err, "Failed to get nuclio function")
		}
		n.logger.DebugWith("Started function", "state", function.Status.State)
		if function.Status.State != functionconfig.FunctionStateReady {
			time.Sleep(3 * time.Second)
		} else {
			break
		}
	}
	return nil
}

func readPlatformConfiguration() (*platformconfig.Config, error) {
	configurationPath := "/etc/nuclio/config/platform/platform.yaml"
	platformConfigurationReader, err := platformconfig.NewReader()
	if err != nil {
		return nil, errors.Wrap(err, "Failed to create platform configuration reader")
	}

	return platformConfigurationReader.ReadFileOrDefault(configurationPath)
}

func getClientConfig(kubeconfigPath string) (*rest.Config, error) {
	if kubeconfigPath != "" {
		return clientcmd.BuildConfigFromFlags("", kubeconfigPath)
	}

	return rest.InClusterConfig()
}
<commit_msg>scale to zero when min replicas is set to zero (#1157)<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/nuclio/nuclio/pkg/errors"
	"github.com/nuclio/nuclio/pkg/functionconfig"
	"github.com/nuclio/nuclio/pkg/loggersink"
	nuclioio_client "github.com/nuclio/nuclio/pkg/platform/kube/client/clientset/versioned"
	// load all sinks
	"github.com/nuclio/nuclio/pkg/platformconfig"
	_ "github.com/nuclio/nuclio/pkg/sinks"

	"github.com/nuclio/logger"
	"github.com/v3io/scaler-types"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

// A plugin for github.com/v3io/scaler, allowing to extend it to scale to zero and from zero function resources in k8s
type NuclioResourceScaler struct {
	logger          logger.Logger
	nuclioClientSet nuclioio_client.Interface
	kubeconfigPath  string
	namespace       string
}

// New is called when plugin loaded on scaler, so it's considered "dead code" for the linter
func New(kubeconfigPath string, namespace string) (scaler_types.ResourceScaler, error) { // nolint: deadcode
	platformConfiguration, err := readPlatformConfiguration()
	if err != nil {
		return nil, errors.Wrap(err, "Failed to read platform configuration")
	}

	// create a root logger
	resourceScalerLogger, err := loggersink.CreateSystemLogger("resource-scaler", platformConfiguration)
	if err != nil {
		fmt.Println(err)
		return nil, errors.Wrap(err, "Failed to create logger")
	}

	restConfig, err := getClientConfig(kubeconfigPath)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to get client configuration")
	}

	nuclioClientSet, err := nuclioio_client.NewForConfig(restConfig)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to create nuclio client set")
	}

	resourceScalerLogger.DebugWith("Initialized resource scaler",
		"platformconfig", platformConfiguration,
		"namespace", namespace,
		"kubeconfigPath", kubeconfigPath)

	return &NuclioResourceScaler{
		logger:          resourceScalerLogger,
		nuclioClientSet: nuclioClientSet,
		kubeconfigPath:  kubeconfigPath,
		namespace:       namespace,
	}, nil
}

func (n *NuclioResourceScaler) SetScale(resource scaler_types.Resource, scale int) error {
	if scale == 0 {
		return n.scaleFunctionToZero(n.namespace, string(resource))
	}
	return n.scaleFunctionFromZero(n.namespace, string(resource))
}

func (n *NuclioResourceScaler) GetResources() ([]scaler_types.Resource, error) {
	functions, err := n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(n.namespace).List(metav1.ListOptions{})
	if err != nil {
		return nil, errors.Wrap(err, "Failed to list functions")
	}

	var functionList []scaler_types.Resource

	// build a list of function names that are potential to be scaled to zero
	for _, function := range functions.Items {

		// don't include functions that aren't in ready state or that min replicas is larger than zero
		if function.Spec.MinReplicas == 0 && function.Status.State == functionconfig.FunctionStateReady {
			functionList = append(functionList, scaler_types.Resource(function.Name))
		}
	}
	return functionList, nil
}

func (n *NuclioResourceScaler) GetConfig() (*scaler_types.ResourceScalerConfig, error) {
	platformConfiguration, err := readPlatformConfiguration()
	if err != nil {
		return nil, errors.Wrap(err, "Failed to read platform configuration")
	}

	scaleInterval, err := time.ParseDuration(platformConfiguration.ScaleToZero.ScalerInterval)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to read scale interval")
	}

	scaleWindow, err := time.ParseDuration(platformConfiguration.ScaleToZero.WindowSize)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to parse scale window")
	}

	pollerInterval, err := time.ParseDuration(platformConfiguration.ScaleToZero.PollerInterval)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to parse poller interval")
	}

	return &scaler_types.ResourceScalerConfig{
		KubeconfigPath: n.kubeconfigPath,
		AutoScalerOptions: scaler_types.AutoScalerOptions{
			Namespace:     n.namespace,
			ScaleInterval: scaleInterval,
			ScaleWindow:   scaleWindow,
			MetricName:    platformConfiguration.ScaleToZero.MetricName,
			Threshold:     0,
		},
		DLXOptions: scaler_types.DLXOptions{
			Namespace:        n.namespace,
			TargetPort:       8080,
			TargetNameHeader: "X-Nuclio-Target",
			TargetPathHeader: "X-Nuclio-Function-Path",
			ListenAddress:    ":8080",
		},
		PollerOptions: scaler_types.PollerOptions{
			MetricInterval: pollerInterval,
			MetricName:     platformConfiguration.ScaleToZero.MetricName,
			Namespace:      n.namespace,
			GroupKind:      "NuclioFunction",
		},
	}, nil
}

func (n *NuclioResourceScaler) scaleFunctionToZero(namespace string, functionName string) error {
	n.logger.DebugWith("Scaling to zero", "functionName", functionName)
	function, err := n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Get(functionName, metav1.GetOptions{})
	if err != nil {
		n.logger.WarnWith("Failed to get nuclio function", "functionName", functionName, "err", err)
		return errors.Wrap(err, "Failed to get function")
	}

	// this has the nice property of disabling hpa as well
	function.Status.State = functionconfig.FunctionStateScaledToZero
	_, err = n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Update(function)
	if err != nil {
		n.logger.WarnWith("Failed to update function", "functionName", functionName, "err", err)
		return errors.Wrap(err, "Failed to update function")
	}
	return nil
}

func (n *NuclioResourceScaler) scaleFunctionFromZero(namespace string, functionName string) error {
	err := n.updateFunctionStatus(namespace, functionName)
	if err != nil {
		return errors.Wrap(err, "Failed to change function status to waitingForResourceConfiguration")
	}
	return n.waitFunctionReadiness(namespace, functionName)
}

func (n *NuclioResourceScaler) updateFunctionStatus(namespace string, functionName string) error {
	function, err := n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Get(functionName, metav1.GetOptions{})
	if err != nil {
		n.logger.WarnWith("Failed to get nuclio function", "functionName", functionName, "err", err)
		return errors.Wrap(err, "Failed to get nuclio function")
	}

	function.Status.State = functionconfig.FunctionStateWaitingForResourceConfiguration
	_, err = n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Update(function)
	if err != nil {
		n.logger.WarnWith("Failed to update function", "functionName", functionName, "err", err)
		return errors.Wrap(err, "Failed to update nuclio function")
	}
	return nil
}

func (n *NuclioResourceScaler) waitFunctionReadiness(namespace string, functionName string) error {
	for {
		function, err := n.nuclioClientSet.NuclioV1beta1().NuclioFunctions(namespace).Get(functionName, metav1.GetOptions{})
		if err != nil {
			n.logger.WarnWith("Failed to get nuclio function", "functionName", functionName, "err", err)
			return errors.Wrap(err, "Failed to get nuclio function")
		}
		n.logger.DebugWith("Started function", "state", function.Status.State)
		if function.Status.State != functionconfig.FunctionStateReady {
			time.Sleep(3 * time.Second)
		} else {
			break
		}
	}
	return nil
}

func readPlatformConfiguration() (*platformconfig.Config, error) {
	configurationPath := "/etc/nuclio/config/platform/platform.yaml"
	platformConfigurationReader, err := platformconfig.NewReader()
	if err != nil {
		return nil, errors.Wrap(err, "Failed to create platform configuration reader")
	}

	return platformConfigurationReader.ReadFileOrDefault(configurationPath)
}

func getClientConfig(kubeconfigPath string) (*rest.Config, error) {
	if kubeconfigPath != "" {
		return clientcmd.BuildConfigFromFlags("", kubeconfigPath)
	}

	return rest.InClusterConfig()
}
<|endoftext|>


<commit_before>/*
 * A name server which sends back the IP address of its client, the
 * recursive resolver. When queried for type TXT, it sends back the text
 * form of the address.  When queried for type A (resp. AAAA), it sends
 * back the IPv4 (resp. v6) address.
 *
 * Similar services: whoami.ultradns.net, whoami.akamai.net. Also (but it
 * is not their normal goal): rs.dns-oarc.net, porttest.dns-oarc.net,
 * amiopen.openresolvers.org.
 *
 * Original version from:
 * Stephane Bortzmeyer <stephane+grong@bortzmeyer.org>
 *
 * Adapted to Go DNS (i.e. completely rewritten)
 * Miek Gieben <miek@miek.nl>
 */

package main

import (
	"flag"
	"fmt"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var (
	printf   *bool
	compress *bool
	tsig     *string
)

const dom = "whoami.miek.nl."

func handleReflect(w dns.ResponseWriter, r *dns.Msg) {
	var (
		v4  bool
		rr  dns.RR
		str string
		a   net.IP
	)
	// TC must be done here
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = *compress
	if ip, ok := w.RemoteAddr().(*net.UDPAddr); ok {
		str = "Port: " + strconv.Itoa(ip.Port) + " (udp)"
		a = ip.IP
		v4 = a.To4() != nil
	}
	if ip, ok := w.RemoteAddr().(*net.TCPAddr); ok {
		str = "Port: " + strconv.Itoa(ip.Port) + " (tcp)"
		a = ip.IP
		v4 = a.To4() != nil
	}

	if v4 {
		rr = new(dns.RR_A)
		rr.(*dns.RR_A).Hdr = dns.RR_Header{Name: dom, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0}
		rr.(*dns.RR_A).A = a.To4()
	} else {
		rr = new(dns.RR_AAAA)
		rr.(*dns.RR_AAAA).Hdr = dns.RR_Header{Name: dom, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 0}
		rr.(*dns.RR_AAAA).AAAA = a
	}

	t := new(dns.RR_TXT)
	t.Hdr = dns.RR_Header{Name: dom, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 0}
	t.Txt = []string{str}

	switch r.Question[0].Qtype {
	case dns.TypeAXFR:
		c := make(chan *dns.XfrToken)
		var e *error
		if err := dns.XfrSend(w, r, c, e); err != nil {
			close(c)
			return
		}
		soa, _ := dns.NewRR(`whoami.miek.nl. IN SOA elektron.atoom.net. miekg.atoom.net. (
			2009032802
			21600
			7200
			604800
			3600)`)
		c <- &dns.XfrToken{RR: []dns.RR{soa, t, rr, soa}}
		close(c)
		w.Hijack()
		// w.Close() // Client closes
		return
	case dns.TypeTXT:
		m.Answer = append(m.Answer, t)
		m.Extra = append(m.Extra, rr)
	default:
		fallthrough
	case dns.TypeAAAA, dns.TypeA:
		m.Answer = append(m.Answer, rr)
		m.Extra = append(m.Extra, t)
	}

	if r.IsTsig() != nil {
		if w.TsigStatus() == nil {
			m.SetTsig(r.Extra[len(r.Extra)-1].(*dns.RR_TSIG).Hdr.Name, dns.HmacMD5, 300, time.Now().Unix())
		} else {
			println("Status", w.TsigStatus().Error())
		}
	}
	if *printf {
		fmt.Printf("%v\n", m.String())
	}
	w.Write(m)
}

func serve(net, name, secret string) {
	switch name {
	case "":
		err := dns.ListenAndServe(":8053", net, nil)
		if err != nil {
			fmt.Printf("Failed to setup the "+net+" server: %s\n", err.Error())
		}
	default:
		server := &dns.Server{Addr: ":8053", Net: net, TsigSecret: map[string]string{name: secret}}
		err := server.ListenAndServe()
		if err != nil {
			fmt.Printf("Failed to setup the "+net+" server: %s\n", err.Error())
		}
	}
}

func main() {
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
	printf = flag.Bool("print", false, "print replies")
	compress = flag.Bool("compress", false, "compress replies")
	tsig = flag.String("tsig", "", "use MD5 hmac tsig: keyname:base64")
	var name, secret string
	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()
	if *tsig != "" {
		a := strings.SplitN(*tsig, ":", 2)
		name, secret = a[0], a[1]
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	dns.HandleFunc("miek.nl.", handleReflect)
	dns.HandleFunc("authors.bind.", dns.HandleAuthors)
	dns.HandleFunc("authors.server.", dns.HandleAuthors)
	dns.HandleFunc("version.bind.", dns.HandleVersion)
	dns.HandleFunc("version.server.", dns.HandleVersion)
	go serve("tcp", name, secret)
	go serve("udp", name, secret)
	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
forever:
	for {
		select {
		case s:=<-sig:
			fmt.Printf("Signal (%d) received, stopping\n", s)
			break forever
		}
	}
}
<commit_msg>documentation<commit_after>// Reflect is a small name server which sends back the IP address of its client, the
// recursive resolver. When queried for type TXT, it sends back the text
// form of the address.  When queried for type A (resp. AAAA), it sends
// back the IPv4 (resp. v6) address.
//
// Similar services: whoami.ultradns.net, whoami.akamai.net. Also (but it
// is not their normal goal): rs.dns-oarc.net, porttest.dns-oarc.net,
// amiopen.openresolvers.org.
//
// Original version is from: Stephane Bortzmeyer <stephane+grong@bortzmeyer.org>.
//
// Adapted to Go (i.e. completely rewritten) by Miek Gieben <miek@miek.nl>.
package main

import (
	"flag"
	"fmt"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var (
	printf   *bool
	compress *bool
	tsig     *string
)

const dom = "whoami.miek.nl."

func handleReflect(w dns.ResponseWriter, r *dns.Msg) {
	var (
		v4  bool
		rr  dns.RR
		str string
		a   net.IP
	)
	// TC must be done here
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = *compress
	if ip, ok := w.RemoteAddr().(*net.UDPAddr); ok {
		str = "Port: " + strconv.Itoa(ip.Port) + " (udp)"
		a = ip.IP
		v4 = a.To4() != nil
	}
	if ip, ok := w.RemoteAddr().(*net.TCPAddr); ok {
		str = "Port: " + strconv.Itoa(ip.Port) + " (tcp)"
		a = ip.IP
		v4 = a.To4() != nil
	}

	if v4 {
		rr = new(dns.RR_A)
		rr.(*dns.RR_A).Hdr = dns.RR_Header{Name: dom, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0}
		rr.(*dns.RR_A).A = a.To4()
	} else {
		rr = new(dns.RR_AAAA)
		rr.(*dns.RR_AAAA).Hdr = dns.RR_Header{Name: dom, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 0}
		rr.(*dns.RR_AAAA).AAAA = a
	}

	t := new(dns.RR_TXT)
	t.Hdr = dns.RR_Header{Name: dom, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 0}
	t.Txt = []string{str}

	switch r.Question[0].Qtype {
	case dns.TypeAXFR:
		c := make(chan *dns.XfrToken)
		var e *error
		if err := dns.XfrSend(w, r, c, e); err != nil {
			close(c)
			return
		}
		soa, _ := dns.NewRR(`whoami.miek.nl. IN SOA elektron.atoom.net. miekg.atoom.net. (
			2009032802
			21600
			7200
			604800
			3600)`)
		c <- &dns.XfrToken{RR: []dns.RR{soa, t, rr, soa}}
		close(c)
		w.Hijack()
		// w.Close() // Client closes
		return
	case dns.TypeTXT:
		m.Answer = append(m.Answer, t)
		m.Extra = append(m.Extra, rr)
	default:
		fallthrough
	case dns.TypeAAAA, dns.TypeA:
		m.Answer = append(m.Answer, rr)
		m.Extra = append(m.Extra, t)
	}

	if r.IsTsig() != nil {
		if w.TsigStatus() == nil {
			m.SetTsig(r.Extra[len(r.Extra)-1].(*dns.RR_TSIG).Hdr.Name, dns.HmacMD5, 300, time.Now().Unix())
		} else {
			println("Status", w.TsigStatus().Error())
		}
	}
	if *printf {
		fmt.Printf("%v\n", m.String())
	}
	w.Write(m)
}

func serve(net, name, secret string) {
	switch name {
	case "":
		err := dns.ListenAndServe(":8053", net, nil)
		if err != nil {
			fmt.Printf("Failed to setup the "+net+" server: %s\n", err.Error())
		}
	default:
		server := &dns.Server{Addr: ":8053", Net: net, TsigSecret: map[string]string{name: secret}}
		err := server.ListenAndServe()
		if err != nil {
			fmt.Printf("Failed to setup the "+net+" server: %s\n", err.Error())
		}
	}
}

func main() {
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")
	printf = flag.Bool("print", false, "print replies")
	compress = flag.Bool("compress", false, "compress replies")
	tsig = flag.String("tsig", "", "use MD5 hmac tsig: keyname:base64")
	var name, secret string
	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()
	if *tsig != "" {
		a := strings.SplitN(*tsig, ":", 2)
		name, secret = a[0], a[1]
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	dns.HandleFunc("miek.nl.", handleReflect)
	dns.HandleFunc("authors.bind.", dns.HandleAuthors)
	dns.HandleFunc("authors.server.", dns.HandleAuthors)
	dns.HandleFunc("version.bind.", dns.HandleVersion)
	dns.HandleFunc("version.server.", dns.HandleVersion)
	go serve("tcp", name, secret)
	go serve("udp", name, secret)
	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
forever:
	for {
		select {
		case s:=<-sig:
			fmt.Printf("Signal (%d) received, stopping\n", s)
			break forever
		}
	}
}
<|endoftext|>


<commit_before>package exec

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"v.io/core/veyron2/vlog"

	"v.io/core/veyron/lib/timekeeper"
)

var (
	ErrAuthTimeout    = errors.New("timeout in auth handshake")
	ErrTimeout        = errors.New("timeout waiting for child")
	ErrSecretTooLarge = errors.New("secret is too large")
)

// A ParentHandle is the Parent process' means of managing a single child.
type ParentHandle struct {
	c           *exec.Cmd
	config      Config
	secret      string
	statusRead  *os.File
	statusWrite *os.File
	tk          timekeeper.TimeKeeper
	waitDone    bool
	waitErr     error
	waitLock    sync.Mutex
	callbackPid int
}

// ParentHandleOpt is an option for NewParentHandle.
type ParentHandleOpt interface {
	// ExecParentHandleOpt is a signature 'dummy' method for the
	// interface.
	ExecParentHandleOpt()
}

// ConfigOpt can be used to seed the parent handle with a
// config to be passed to the child.
type ConfigOpt struct {
	Config
}

// ExecParentHandleOpt makes ConfigOpt an instance of
// ParentHandleOpt.
func (ConfigOpt) ExecParentHandleOpt() {}

// SecretOpt can be used to seed the parent handle with a custom secret.
type SecretOpt string

// ExecParentHandleOpt makes SecretOpt an instance of ParentHandleOpt.
func (SecretOpt) ExecParentHandleOpt() {}

// TimeKeeperOpt can be used to seed the parent handle with a custom timekeeper.
type TimeKeeperOpt struct {
	timekeeper.TimeKeeper
}

// ExecParentHandleOpt makes TimeKeeperOpt an instance of ParentHandleOpt.
func (TimeKeeperOpt) ExecParentHandleOpt() {}

// NewParentHandle creates a ParentHandle for the child process represented by
// an instance of exec.Cmd.
func NewParentHandle(c *exec.Cmd, opts ...ParentHandleOpt) *ParentHandle {
	cfg, secret := NewConfig(), ""
	tk := timekeeper.RealTime()
	for _, opt := range opts {
		switch v := opt.(type) {
		case ConfigOpt:
			cfg = v
		case SecretOpt:
			secret = string(v)
		case TimeKeeperOpt:
			tk = v
		default:
			vlog.Errorf("Unrecognized parent option: %v", v)
		}
	}
	return &ParentHandle{
		c:      c,
		config: cfg,
		secret: secret,
		tk:     tk,
	}
}

// Start starts the child process, sharing a secret with it and
// setting up a communication channel over which to read its status.
func (p *ParentHandle) Start() error {
	// Make sure that there are no instances of the VersionVariable
	// already in the environment (which can happen when a subprocess
	// creates a subprocess etc)
	nenv := make([]string, 0, len(p.c.Env)+1)
	for _, e := range p.c.Env {
		if strings.HasPrefix(e, VersionVariable+"=") {
			continue
		}
		nenv = append(nenv, e)
	}
	p.c.Env = append(nenv, VersionVariable+"="+version1)

	// Create anonymous pipe for communicating data between the child
	// and the parent.
	// TODO(caprita): As per ribrdb@, Go's exec does not prune the set
	// of file descriptors passed down to the child process, and hence
	// a child may get access to the files meant for another child.
	// Do we need to ensure only one thread is allowed to create these
	// pipes at any time?
	dataRead, dataWrite, err := os.Pipe()
	if err != nil {
		return err
	}
	defer dataRead.Close()
	defer dataWrite.Close()
	statusRead, statusWrite, err := os.Pipe()
	if err != nil {
		return err
	}
	p.statusRead = statusRead
	p.statusWrite = statusWrite
	// Add the parent-child pipes to cmd.ExtraFiles, offsetting all
	// existing file descriptors accordingly.
	extraFiles := make([]*os.File, len(p.c.ExtraFiles)+2)
	extraFiles[0] = dataRead
	extraFiles[1] = statusWrite
	for i, _ := range p.c.ExtraFiles {
		extraFiles[i+2] = p.c.ExtraFiles[i]
	}
	p.c.ExtraFiles = extraFiles
	// Start the child process.
	if err := p.c.Start(); err != nil {
		p.statusWrite.Close()
		p.statusRead.Close()
		return err
	}
	// Pass data to the child using a pipe.
	serializedConfig, err := p.config.Serialize()
	if err != nil {
		return err
	}
	if err := encodeString(dataWrite, serializedConfig); err != nil {
		p.statusWrite.Close()
		p.statusRead.Close()
		return err
	}
	if err := encodeString(dataWrite, p.secret); err != nil {
		p.statusWrite.Close()
		p.statusRead.Close()
		return err
	}
	return nil
}

// copy is like io.Copy, but it also treats the receipt of the special eofChar
// byte to mean io.EOF.
func copy(w io.Writer, r io.Reader) (err error) {
	buf := make([]byte, 1024)
	for {
		nRead, errRead := r.Read(buf)
		if nRead > 0 {
			if eofCharIndex := bytes.IndexByte(buf[:nRead], eofChar); eofCharIndex != -1 {
				nRead = eofCharIndex
				errRead = io.EOF
			}
			nWrite, errWrite := w.Write(buf[:nRead])
			if errWrite != nil {
				err = errWrite
				break
			}
			if nRead != nWrite {
				err = io.ErrShortWrite
				break
			}
		}
		if errRead == io.EOF {
			break
		}
		if errRead != nil {
			err = errRead
			break
		}
	}
	return
}

func waitForStatus(c chan interface{}, r *os.File) {
	var readBytes bytes.Buffer
	err := copy(&readBytes, r)
	r.Close()
	if err != nil {
		c <- err
	} else {
		c <- readBytes.String()
	}
	close(c)
}

// WaitForReady will wait for the child process to become ready.
func (p *ParentHandle) WaitForReady(timeout time.Duration) error {
	// An invariant of WaitForReady is that both statusWrite and statusRead
	// get closed before WaitForStatus returns (statusRead gets closed by
	// waitForStatus).
	defer p.statusWrite.Close()
	c := make(chan interface{}, 1)
	go waitForStatus(c, p.statusRead)
	// TODO(caprita): This can be simplified further by doing the reading
	// from the status pipe here, and instead moving the timeout listener to
	// a separate goroutine.
	select {
	case msg := <-c:
		switch m := msg.(type) {
		case error:
			return m
		case string:
			if strings.HasPrefix(m, readyStatus) {
				pid, err := strconv.Atoi(m[len(readyStatus):])
				if err != nil {
					return err
				}
				p.callbackPid = pid
				return nil
			}
			if strings.HasPrefix(m, failedStatus) {
				return fmt.Errorf("%s", strings.TrimPrefix(m, failedStatus))
			}
			return fmt.Errorf("unrecognised status from subprocess: %q", m)
		default:
			return fmt.Errorf("unexpected type %T", m)
		}
	case <-p.tk.After(timeout):
		vlog.Errorf("Timed out waiting for child status")
		// By writing the special eofChar byte to the pipe, we ensure
		// that waitForStatus returns: the copy function treats eofChar
		// to indicate end of read input.  Note, copy could have
		// finished for other reasons already (receipt of eofChar from
		// the child process).  Note, closing the pipe from the child
		// (explicitly or due to crash) would NOT cause copy to read
		// io.EOF, since we keep the statusWrite open in the parent.
		// Hence, a child crash will eventually trigger this timeout.
		p.statusWrite.Write([]byte{eofChar})
		// Before returning, waitForStatus will close r, and then close
		// c.  Waiting on c ensures that r.Close() in waitForStatus
		// already executed.
		<-c
		return ErrTimeout
	}
	panic("unreachable")
}

// wait performs the Wait on the underlying command under lock, and only once
// (subsequent wait calls block until the Wait is finished).  It's ok to call
// wait multiple times, and in parallel.  The error from the initial Wait is
// cached and returned for all subsequent calls.
func (p *ParentHandle) wait() error {
	p.waitLock.Lock()
	defer p.waitLock.Unlock()
	if p.waitDone {
		return p.waitErr
	}
	p.waitErr = p.c.Wait()
	p.waitDone = true
	return p.waitErr
}

// Wait will wait for the child process to terminate of its own accord.
// It returns nil if the process exited cleanly with an exit status of 0,
// any other exit code or error will result in an appropriate error return
func (p *ParentHandle) Wait(timeout time.Duration) error {
	c := make(chan error, 1)
	go func() {
		c <- p.wait()
		close(c)
	}()
	// If timeout is zero time.After will panic; we handle zero specially
	// to mean infinite timeout.
	if timeout > 0 {
		select {
		case <-p.tk.After(timeout):
			return ErrTimeout
		case err := <-c:
			return err
		}
	} else {
		return <-c
	}
	panic("unreachable")
}

// Pid returns the pid of the child, 0 if the child process doesn't exist
func (p *ParentHandle) Pid() int {
	if p.c.Process != nil {
		return p.c.Process.Pid
	}
	return 0
}

// ChildPid returns the pid of a child process as reported by its status
// callback.
func (p *ParentHandle) ChildPid() int {
	return p.callbackPid
}

// Exists returns true if the child process exists and can be signal'ed
func (p *ParentHandle) Exists() bool {
	if p.c.Process != nil {
		return syscall.Kill(p.c.Process.Pid, 0) == nil
	}
	return false
}

// Kill kills the child process.
func (p *ParentHandle) Kill() error {
	return p.c.Process.Kill()
}

// Signal sends the given signal to the child process.
func (p *ParentHandle) Signal(sig syscall.Signal) error {
	return syscall.Kill(p.c.Process.Pid, sig)
}

// Clean will clean up state, including killing the child process.
func (p *ParentHandle) Clean() error {
	if err := p.Kill(); err != nil {
		return err
	}
	return p.wait()
}

func encodeString(w io.Writer, data string) error {
	l := len(data)
	if err := binary.Write(w, binary.BigEndian, int64(l)); err != nil {
		return err
	}
	if n, err := w.Write([]byte(data)); err != nil || n != l {
		if err != nil {
			return err
		} else {
			return errors.New("partial write")
		}
	}
	return nil
}
<commit_msg>veyron/lib/exec: Fix panic in handle.Clean()<commit_after>package exec

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"v.io/core/veyron2/vlog"

	"v.io/core/veyron/lib/timekeeper"
)

var (
	ErrAuthTimeout    = errors.New("timeout in auth handshake")
	ErrTimeout        = errors.New("timeout waiting for child")
	ErrSecretTooLarge = errors.New("secret is too large")
)

// A ParentHandle is the Parent process' means of managing a single child.
type ParentHandle struct {
	c           *exec.Cmd
	config      Config
	secret      string
	statusRead  *os.File
	statusWrite *os.File
	tk          timekeeper.TimeKeeper
	waitDone    bool
	waitErr     error
	waitLock    sync.Mutex
	callbackPid int
}

// ParentHandleOpt is an option for NewParentHandle.
type ParentHandleOpt interface {
	// ExecParentHandleOpt is a signature 'dummy' method for the
	// interface.
	ExecParentHandleOpt()
}

// ConfigOpt can be used to seed the parent handle with a
// config to be passed to the child.
type ConfigOpt struct {
	Config
}

// ExecParentHandleOpt makes ConfigOpt an instance of
// ParentHandleOpt.
func (ConfigOpt) ExecParentHandleOpt() {}

// SecretOpt can be used to seed the parent handle with a custom secret.
type SecretOpt string

// ExecParentHandleOpt makes SecretOpt an instance of ParentHandleOpt.
func (SecretOpt) ExecParentHandleOpt() {}

// TimeKeeperOpt can be used to seed the parent handle with a custom timekeeper.
type TimeKeeperOpt struct {
	timekeeper.TimeKeeper
}

// ExecParentHandleOpt makes TimeKeeperOpt an instance of ParentHandleOpt.
func (TimeKeeperOpt) ExecParentHandleOpt() {}

// NewParentHandle creates a ParentHandle for the child process represented by
// an instance of exec.Cmd.
func NewParentHandle(c *exec.Cmd, opts ...ParentHandleOpt) *ParentHandle {
	cfg, secret := NewConfig(), ""
	tk := timekeeper.RealTime()
	for _, opt := range opts {
		switch v := opt.(type) {
		case ConfigOpt:
			cfg = v
		case SecretOpt:
			secret = string(v)
		case TimeKeeperOpt:
			tk = v
		default:
			vlog.Errorf("Unrecognized parent option: %v", v)
		}
	}
	return &ParentHandle{
		c:      c,
		config: cfg,
		secret: secret,
		tk:     tk,
	}
}

// Start starts the child process, sharing a secret with it and
// setting up a communication channel over which to read its status.
func (p *ParentHandle) Start() error {
	// Make sure that there are no instances of the VersionVariable
	// already in the environment (which can happen when a subprocess
	// creates a subprocess etc)
	nenv := make([]string, 0, len(p.c.Env)+1)
	for _, e := range p.c.Env {
		if strings.HasPrefix(e, VersionVariable+"=") {
			continue
		}
		nenv = append(nenv, e)
	}
	p.c.Env = append(nenv, VersionVariable+"="+version1)

	// Create anonymous pipe for communicating data between the child
	// and the parent.
	// TODO(caprita): As per ribrdb@, Go's exec does not prune the set
	// of file descriptors passed down to the child process, and hence
	// a child may get access to the files meant for another child.
	// Do we need to ensure only one thread is allowed to create these
	// pipes at any time?
	dataRead, dataWrite, err := os.Pipe()
	if err != nil {
		return err
	}
	defer dataRead.Close()
	defer dataWrite.Close()
	statusRead, statusWrite, err := os.Pipe()
	if err != nil {
		return err
	}
	p.statusRead = statusRead
	p.statusWrite = statusWrite
	// Add the parent-child pipes to cmd.ExtraFiles, offsetting all
	// existing file descriptors accordingly.
	extraFiles := make([]*os.File, len(p.c.ExtraFiles)+2)
	extraFiles[0] = dataRead
	extraFiles[1] = statusWrite
	for i, _ := range p.c.ExtraFiles {
		extraFiles[i+2] = p.c.ExtraFiles[i]
	}
	p.c.ExtraFiles = extraFiles
	// Start the child process.
	if err := p.c.Start(); err != nil {
		p.statusWrite.Close()
		p.statusRead.Close()
		return err
	}
	// Pass data to the child using a pipe.
	serializedConfig, err := p.config.Serialize()
	if err != nil {
		return err
	}
	if err := encodeString(dataWrite, serializedConfig); err != nil {
		p.statusWrite.Close()
		p.statusRead.Close()
		return err
	}
	if err := encodeString(dataWrite, p.secret); err != nil {
		p.statusWrite.Close()
		p.statusRead.Close()
		return err
	}
	return nil
}

// copy is like io.Copy, but it also treats the receipt of the special eofChar
// byte to mean io.EOF.
func copy(w io.Writer, r io.Reader) (err error) {
	buf := make([]byte, 1024)
	for {
		nRead, errRead := r.Read(buf)
		if nRead > 0 {
			if eofCharIndex := bytes.IndexByte(buf[:nRead], eofChar); eofCharIndex != -1 {
				nRead = eofCharIndex
				errRead = io.EOF
			}
			nWrite, errWrite := w.Write(buf[:nRead])
			if errWrite != nil {
				err = errWrite
				break
			}
			if nRead != nWrite {
				err = io.ErrShortWrite
				break
			}
		}
		if errRead == io.EOF {
			break
		}
		if errRead != nil {
			err = errRead
			break
		}
	}
	return
}

func waitForStatus(c chan interface{}, r *os.File) {
	var readBytes bytes.Buffer
	err := copy(&readBytes, r)
	r.Close()
	if err != nil {
		c <- err
	} else {
		c <- readBytes.String()
	}
	close(c)
}

// WaitForReady will wait for the child process to become ready.
func (p *ParentHandle) WaitForReady(timeout time.Duration) error {
	// An invariant of WaitForReady is that both statusWrite and statusRead
	// get closed before WaitForStatus returns (statusRead gets closed by
	// waitForStatus).
	defer p.statusWrite.Close()
	c := make(chan interface{}, 1)
	go waitForStatus(c, p.statusRead)
	// TODO(caprita): This can be simplified further by doing the reading
	// from the status pipe here, and instead moving the timeout listener to
	// a separate goroutine.
	select {
	case msg := <-c:
		switch m := msg.(type) {
		case error:
			return m
		case string:
			if strings.HasPrefix(m, readyStatus) {
				pid, err := strconv.Atoi(m[len(readyStatus):])
				if err != nil {
					return err
				}
				p.callbackPid = pid
				return nil
			}
			if strings.HasPrefix(m, failedStatus) {
				return fmt.Errorf("%s", strings.TrimPrefix(m, failedStatus))
			}
			return fmt.Errorf("unrecognised status from subprocess: %q", m)
		default:
			return fmt.Errorf("unexpected type %T", m)
		}
	case <-p.tk.After(timeout):
		vlog.Errorf("Timed out waiting for child status")
		// By writing the special eofChar byte to the pipe, we ensure
		// that waitForStatus returns: the copy function treats eofChar
		// to indicate end of read input.  Note, copy could have
		// finished for other reasons already (receipt of eofChar from
		// the child process).  Note, closing the pipe from the child
		// (explicitly or due to crash) would NOT cause copy to read
		// io.EOF, since we keep the statusWrite open in the parent.
		// Hence, a child crash will eventually trigger this timeout.
		p.statusWrite.Write([]byte{eofChar})
		// Before returning, waitForStatus will close r, and then close
		// c.  Waiting on c ensures that r.Close() in waitForStatus
		// already executed.
		<-c
		return ErrTimeout
	}
	panic("unreachable")
}

// wait performs the Wait on the underlying command under lock, and only once
// (subsequent wait calls block until the Wait is finished).  It's ok to call
// wait multiple times, and in parallel.  The error from the initial Wait is
// cached and returned for all subsequent calls.
func (p *ParentHandle) wait() error {
	p.waitLock.Lock()
	defer p.waitLock.Unlock()
	if p.waitDone {
		return p.waitErr
	}
	p.waitErr = p.c.Wait()
	p.waitDone = true
	return p.waitErr
}

// Wait will wait for the child process to terminate of its own accord.
// It returns nil if the process exited cleanly with an exit status of 0,
// any other exit code or error will result in an appropriate error return
func (p *ParentHandle) Wait(timeout time.Duration) error {
	c := make(chan error, 1)
	go func() {
		c <- p.wait()
		close(c)
	}()
	// If timeout is zero time.After will panic; we handle zero specially
	// to mean infinite timeout.
	if timeout > 0 {
		select {
		case <-p.tk.After(timeout):
			return ErrTimeout
		case err := <-c:
			return err
		}
	} else {
		return <-c
	}
	panic("unreachable")
}

// Pid returns the pid of the child, 0 if the child process doesn't exist
func (p *ParentHandle) Pid() int {
	if p.c.Process != nil {
		return p.c.Process.Pid
	}
	return 0
}

// ChildPid returns the pid of a child process as reported by its status
// callback.
func (p *ParentHandle) ChildPid() int {
	return p.callbackPid
}

// Exists returns true if the child process exists and can be signal'ed
func (p *ParentHandle) Exists() bool {
	if p.c.Process != nil {
		return syscall.Kill(p.c.Process.Pid, 0) == nil
	}
	return false
}

// Kill kills the child process.
func (p *ParentHandle) Kill() error {
	if p.c.Process == nil {
		return errors.New("no such process")
	}
	return p.c.Process.Kill()
}

// Signal sends the given signal to the child process.
func (p *ParentHandle) Signal(sig syscall.Signal) error {
	if p.c.Process == nil {
		return errors.New("no such process")
	}
	return syscall.Kill(p.c.Process.Pid, sig)
}

// Clean will clean up state, including killing the child process.
func (p *ParentHandle) Clean() error {
	if err := p.Kill(); err != nil {
		return err
	}
	return p.wait()
}

func encodeString(w io.Writer, data string) error {
	l := len(data)
	if err := binary.Write(w, binary.BigEndian, int64(l)); err != nil {
		return err
	}
	if n, err := w.Write([]byte(data)); err != nil || n != l {
		if err != nil {
			return err
		} else {
			return errors.New("partial write")
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/1lann/beacon/chat"
	"github.com/1lann/beacon/handler"
	"io"
	"log"
	"net"
	"strings"
	"sync"
)

type state int

var currentState state

const (
	stateInitializing = iota
	stateStopped
	stateOff
	stateShutdown
	stateSnapshot
	stateDestroy
	stateStarted
	stateStarting
	stateUnavailable
)

const prefixText = chat.Aqua + "[chuie.io] " + chat.White
const headerText = chat.Aqua + "-- [ chuie.io ] --\n\n" + chat.White

func (s state) String() string {
	switch s {
	case stateInitializing:
		return "Intializing"
	case stateStopped:
		return "Stopped"
	case stateOff:
		return "Off"
	case stateSnapshot:
		return "Snapshot"
	case stateShutdown:
		return "Shutdown"
	case stateDestroy:
		return "Destroy"
	case stateStarted:
		return "Started"
	case stateStarting:
		return "Starting"
	}

	return "Unknown"
}

var connectMessage string

func setState(s state) {
	// if time.Now().Sub(lastSet) < time.Duration(time.Second*2) {
	// 	log.Println("Set state cancelled due to being too fast")
	// }

	// lastSet := time.Now()

	if s != stateStarted && currentState == stateStarted {
		stopForwarder()
		go startBeacon()
	}

	switch s {
	case stateInitializing:
		connectMessage = "Sorry, the server is not ready take requests yet!\n\n" +
			"Try connecting again in a few seconds."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText +
			chat.Yellow + "Initializing..."
		handler.CurrentStatus.ShowConnection = false
	case stateStopped:
		connectMessage = "Sorry, the server is intentionally down.\n\n" +
			"Contact the server owner for more information."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText +
			chat.Red + "Intentionally down"
		handler.CurrentStatus.ShowConnection = false
	case stateShutdown, stateSnapshot, stateDestroy:
		connectMessage = "Sorry, the server is currently shutting down.\n" +
			"You may start it again when it is completely turned off.\n\n" +
			"Try connecting again in a few minutes."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText + chat.Yellow +
			"Shutting down..."
		handler.CurrentStatus.ShowConnection = false
	case stateOff:
		handler.OnConnect = onConnectIdle
		handler.CurrentStatus.Message = prefixText + "Shutdown. Connect to start."
		handler.CurrentStatus.ShowConnection = true
	case stateStarting:
		connectMessage = "Sorry, the server is still starting up.\n\n" +
			"Try connecting again in a few minutes."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText + chat.LightGreen +
			"Starting up..."
		handler.CurrentStatus.ShowConnection = false
	case stateUnavailable:
		connectMessage = "The server is unavailable due to an error.\n\n" +
			"Contact the server owner for help."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText + chat.Red + "Unavailable"
		handler.CurrentStatus.ShowConnection = false
	case stateStarted:
		if currentState != stateStarted {
			stopBeacon()
			go startForwarder()
		}
	}

	currentState = s
}

func onConnectMessage(player *handler.Player) string {
	return headerText + connectMessage
}

func onConnectIdle(player *handler.Player) string {
	log.Println(player.Username + " has started the server.")
	go restoreServer()
	return headerText +
		"The server is now starting. Come back in about 5 minutes."
}

var forwardListener net.Listener
var forwardAddr string
var connectionLock *sync.Mutex = &sync.Mutex{}

func startBeacon() {
	connectionLock.Lock()
	defer connectionLock.Unlock()

	log.Println("[Beacon] Started.")
	err := handler.Listen("25565")
	if err != nil {
		log.Fatal(err)
	}

	log.Println("[Beacon] Stopped.")
}

func stopBeacon() {
	handler.Stop()
}

func startForwarder() {
	connectionLock.Lock()
	defer connectionLock.Unlock()

	var err error
	forwardListener, err = net.Listen("tcp", ":25565")
	if err != nil {
		log.Println("[Forwarder] Failed to listen:", err)
		return
	}

	log.Println("[Forwarder] Started.")

	for {
		conn, err := forwardListener.Accept()
		if err != nil {
			if strings.Contains(err.Error(),
				"use of closed network connection") || err == io.EOF {
				log.Println("[Forwarder] Stopped.")
				return
			}

			log.Println("[Forwarder] Failed to accept:", err)
			return
		}

		go forwardConnection(conn)
	}
}

func forwardConnection(localConn net.Conn) {
	defer localConn.Close()

	if forwardAddr == "" {
		return
	}

	remoteConn, err := net.Dial("tcp", forwardAddr+":25565")
	if err != nil {
		log.Println("[Forwarder] Failed to connect to remote:", err)
		return
	}

	defer remoteConn.Close()

	connChannel := make(chan bool)

	go func() {
		_, _ = io.Copy(remoteConn, localConn)
		connChannel <- true
	}()

	go func() {
		_, _ = io.Copy(localConn, remoteConn)
		connChannel <- true
	}()

	<-connChannel
}

func stopForwarder() {
	_ = forwardListener.Close()
}
<commit_msg>Add color to off message<commit_after>package main

import (
	"github.com/1lann/beacon/chat"
	"github.com/1lann/beacon/handler"
	"io"
	"log"
	"net"
	"strings"
	"sync"
)

type state int

var currentState state

const (
	stateInitializing = iota
	stateStopped
	stateOff
	stateShutdown
	stateSnapshot
	stateDestroy
	stateStarted
	stateStarting
	stateUnavailable
)

const prefixText = chat.Aqua + "[chuie.io] " + chat.White
const headerText = chat.Aqua + "-- [ chuie.io ] --\n\n" + chat.White

func (s state) String() string {
	switch s {
	case stateInitializing:
		return "Intializing"
	case stateStopped:
		return "Stopped"
	case stateOff:
		return "Off"
	case stateSnapshot:
		return "Snapshot"
	case stateShutdown:
		return "Shutdown"
	case stateDestroy:
		return "Destroy"
	case stateStarted:
		return "Started"
	case stateStarting:
		return "Starting"
	}

	return "Unknown"
}

var connectMessage string

func setState(s state) {
	// if time.Now().Sub(lastSet) < time.Duration(time.Second*2) {
	// 	log.Println("Set state cancelled due to being too fast")
	// }

	// lastSet := time.Now()

	if s != stateStarted && currentState == stateStarted {
		stopForwarder()
		go startBeacon()
	}

	switch s {
	case stateInitializing:
		connectMessage = "Sorry, the server is not ready take requests yet!\n\n" +
			"Try connecting again in a few seconds."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText +
			chat.Yellow + "Initializing..."
		handler.CurrentStatus.ShowConnection = false
	case stateStopped:
		connectMessage = "Sorry, the server is intentionally down.\n\n" +
			"Contact the server owner for more information."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText +
			chat.Red + "Intentionally down"
		handler.CurrentStatus.ShowConnection = false
	case stateShutdown, stateSnapshot, stateDestroy:
		connectMessage = "Sorry, the server is currently shutting down.\n" +
			"You may start it again when it is completely turned off.\n\n" +
			"Try connecting again in a few minutes."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText + chat.Yellow +
			"Shutting down..."
		handler.CurrentStatus.ShowConnection = false
	case stateOff:
		handler.OnConnect = onConnectIdle
		handler.CurrentStatus.Message = prefixText + chat.Gold +
			"Powered off. Connect to start."
		handler.CurrentStatus.ShowConnection = true
	case stateStarting:
		connectMessage = "Sorry, the server is still starting up.\n\n" +
			"Try connecting again in a few minutes."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText + chat.LightGreen +
			"Starting up..."
		handler.CurrentStatus.ShowConnection = false
	case stateUnavailable:
		connectMessage = "The server is unavailable due to an error.\n\n" +
			"Contact the server owner for help."
		handler.OnConnect = onConnectMessage
		handler.CurrentStatus.Message = prefixText + chat.Red + "Unavailable"
		handler.CurrentStatus.ShowConnection = false
	case stateStarted:
		if currentState != stateStarted {
			stopBeacon()
			go startForwarder()
		}
	}

	currentState = s
}

func onConnectMessage(player *handler.Player) string {
	return headerText + connectMessage
}

func onConnectIdle(player *handler.Player) string {
	log.Println(player.Username + " has started the server.")
	go restoreServer()
	return headerText +
		"The server is now starting. Come back in about 5 minutes."
}

var forwardListener net.Listener
var forwardAddr string
var connectionLock *sync.Mutex = &sync.Mutex{}

func startBeacon() {
	connectionLock.Lock()
	defer connectionLock.Unlock()

	log.Println("[Beacon] Started.")
	err := handler.Listen("25565")
	if err != nil {
		log.Fatal(err)
	}

	log.Println("[Beacon] Stopped.")
}

func stopBeacon() {
	handler.Stop()
}

func startForwarder() {
	connectionLock.Lock()
	defer connectionLock.Unlock()

	var err error
	forwardListener, err = net.Listen("tcp", ":25565")
	if err != nil {
		log.Println("[Forwarder] Failed to listen:", err)
		return
	}

	log.Println("[Forwarder] Started.")

	for {
		conn, err := forwardListener.Accept()
		if err != nil {
			if strings.Contains(err.Error(),
				"use of closed network connection") || err == io.EOF {
				log.Println("[Forwarder] Stopped.")
				return
			}

			log.Println("[Forwarder] Failed to accept:", err)
			return
		}

		go forwardConnection(conn)
	}
}

func forwardConnection(localConn net.Conn) {
	defer localConn.Close()

	if forwardAddr == "" {
		return
	}

	remoteConn, err := net.Dial("tcp", forwardAddr+":25565")
	if err != nil {
		log.Println("[Forwarder] Failed to connect to remote:", err)
		return
	}

	defer remoteConn.Close()

	connChannel := make(chan bool)

	go func() {
		_, _ = io.Copy(remoteConn, localConn)
		connChannel <- true
	}()

	go func() {
		_, _ = io.Copy(localConn, remoteConn)
		connChannel <- true
	}()

	<-connChannel
}

func stopForwarder() {
	_ = forwardListener.Close()
}
<|endoftext|>