<commit_before>package gateway

import (
	"errors"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

// rpcID is an 8-byte signature that is added to all RPCs to tell the gatway
// what to do with the RPC. Empty elements of rpcID will be encoded as spaces.
type rpcID [8]byte

// String returns a string representation of an rpcID.
func (id rpcID) String() string {
	for i := range id {
		if id[i] == 0 {
			id[i] = ' '
		}
	}
	return string(id[:])
}

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) (id rpcID) {
	copy(id[:], name)
	return
}

// managedRPC calls an RPC on the given address. managedRPC cannot be called on
// an address that the Gateway is not connected to.
func (g *Gateway) managedRPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	g.mu.RLock()
	peer, ok := g.peers[addr]
	g.mu.RUnlock()
	if !ok {
		return errors.New("can't call RPC on unconnected peer " + string(addr))
	}

	conn, err := peer.open()
	if err != nil {
		return err
	}
	defer conn.Close()

	// write header
	if err := encoding.WriteObject(conn, handlerName(name)); err != nil {
		return err
	}
	// call fn
	return fn(conn)
}

// RPC calls an RPC on the given address. RPC cannot be called on an address
// that the Gateway is not connected to.
func (g *Gateway) RPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	if err := g.threads.Add(); err != nil {
		return err
	}
	defer g.threads.Done()
	return g.managedRPC(addr, name, fn)
}

// RegisterRPC registers an RPCFunc as a handler for a given identifier. To
// call an RPC, use gateway.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase. The first 8
// characters of an identifier should be unique, as the identifier used
// internally is truncated to 8 bytes.
func (g *Gateway) RegisterRPC(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; ok {
		build.Critical("RPC already registered: " + name)
	}
	g.handlers[handlerName(name)] = fn
}

// UnregisterRPC unregisters an RPC and removes the corresponding RPCFunc from
// g.handlers. Future calls to the RPC by peers will fail.
func (g *Gateway) UnregisterRPC(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; !ok {
		build.Critical("RPC not registered: " + name)
	}
	delete(g.handlers, handlerName(name))
}

// RegisterConnectCall registers a name and RPCFunc to be called on a peer
// upon connecting.
func (g *Gateway) RegisterConnectCall(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; ok {
		build.Critical("ConnectCall already registered: " + name)
	}
	g.initRPCs[name] = fn
}

// UnregisterConnectCall unregisters an on-connect call and removes the
// corresponding RPCFunc from g.initRPCs. Future connections to peers will not
// trigger the RPC to be called on them.
func (g *Gateway) UnregisterConnectCall(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; !ok {
		build.Critical("ConnectCall not registered: " + name)
	}
	delete(g.initRPCs, name)
}

// threadedListenPeer listens for new streams on a peer connection and serves them via
// threadedHandleConn.
func (g *Gateway) threadedListenPeer(p *peer) {
	// threadedListenPeer registers to the peerTG instead of the primary thread
	// group because peer connections can be lifetime in length, but can also
	// be short-lived. The fact that they can be lifetime means that they can't
	// call threads.Add as they will block calls to threads.Flush. The fact
	// that they can be short-lived means that threads.OnStop is not a good
	// tool for closing out the threads. Instead, they register to peerTG,
	// which is cleanly closed upon gateway shutdown but will not block any
	// calls to threads.Flush()
	if g.peerTG.Add() != nil {
		return
	}
	defer g.peerTG.Done()

	// Spin up a goroutine to listen for a shutdown signal from both the peer
	// and from the gateway. In the event of either, close the muxado session.
	connClosedChan := make(chan struct{})
	peerCloseChan := make(chan struct{})
	go func() {
		// Signal that the muxado session has been successfully closed, and
		// that this goroutine has terminated.
		defer close(connClosedChan)

		// Listen for a stop signal.
		select {
		case <-g.threads.StopChan():
		case <-peerCloseChan:
		}

		// Can't call Disconnect because it could return sync.ErrStopped.
		g.mu.Lock()
		delete(g.peers, p.NetAddress)
		g.mu.Unlock()
		if err := p.sess.Close(); err != nil {
			g.log.Debugf("WARN: error disconnecting from peer %q: %v", p.NetAddress, err)
		}
	}()

	for {
		conn, err := p.accept()
		if err != nil {
			g.log.Debugln("Peer connection closed:", p.NetAddress)
			break
		}

		// it is the handler's responsibility to close the connection
		go g.threadedHandleConn(conn)
	}
	// Signal that the goroutine can shutdown.
	close(peerCloseChan)
	// Wait for confirmation that the goroutine has shut down before returning
	// and releasing the threadgroup registration.
	<-connClosedChan
}

// threadedHandleConn reads header data from a connection, then routes it to the
// appropriate handler for further processing.
func (g *Gateway) threadedHandleConn(conn modules.PeerConn) {
	defer conn.Close()
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	var id rpcID
	if err := encoding.ReadObject(conn, &id, 8); err != nil {
		return
	}
	// call registered handler for this ID
	g.mu.RLock()
	fn, ok := g.handlers[id]
	g.mu.RUnlock()
	if !ok {
		g.log.Debugf("WARN: incoming conn %v requested unknown RPC \"%v\"", conn.RPCAddr(), id)
		return
	}
	g.log.Debugf("INFO: incoming conn %v requested RPC \"%v\"", conn.RPCAddr(), id)

	// call fn
	err := fn(conn)
	// don't log benign errors
	if err == modules.ErrDuplicateTransactionSet || err == modules.ErrBlockKnown {
		err = nil
	}
	if err != nil {
		g.log.Debugf("WARN: incoming RPC \"%v\" from conn %v failed: %v", id, conn.RPCAddr(), err)
	}
}

// Broadcast calls an RPC on all of the specified peers. The calls are run in
// parallel. Broadcasts are restricted to "one-way" RPCs, which simply write an
// object and disconnect. This is why Broadcast takes an interface{} instead of
// an RPCFunc.
func (g *Gateway) Broadcast(name string, obj interface{}, peers []modules.Peer) {
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	g.log.Debugf("INFO: broadcasting RPC %q to %v peers", name, len(peers))

	// only encode obj once, instead of using WriteObject
	enc := encoding.Marshal(obj)
	fn := func(conn modules.PeerConn) error {
		return encoding.WritePrefix(conn, enc)
	}

	var wg sync.WaitGroup
	for _, p := range peers {
		wg.Add(1)
		go func(addr modules.NetAddress) {
			defer wg.Done()
			err := g.managedRPC(addr, name, fn)
			if err != nil {
				g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed (attempting again in 10 seconds): %v", name, addr, err)
				// try one more time before giving up
				select {
				case <-time.After(10 * time.Second):
				case <-g.threads.StopChan():
					return
				}
				err := g.managedRPC(addr, name, fn)
				if err != nil {
					g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed twice: %v", name, addr, err)
				}
			}
		}(p.NetAddress)
	}
	wg.Wait()
}
<commit_msg>better documentation on rpcID String() encoding<commit_after>package gateway

import (
	"errors"
	"sync"
	"time"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

// rpcID is an 8-byte signature that is added to all RPCs to tell the gatway
// what to do with the RPC.
type rpcID [8]byte

// String returns a string representation of an rpcID. Empty elements of rpcID
// will be encoded as spaces.
func (id rpcID) String() string {
	for i := range id {
		if id[i] == 0 {
			id[i] = ' '
		}
	}
	return string(id[:])
}

// handlerName truncates a string to 8 bytes. If len(name) < 8, the remaining
// bytes are 0. A handlerName is specified at the beginning of each network
// call, indicating which function should handle the connection.
func handlerName(name string) (id rpcID) {
	copy(id[:], name)
	return
}

// managedRPC calls an RPC on the given address. managedRPC cannot be called on
// an address that the Gateway is not connected to.
func (g *Gateway) managedRPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	g.mu.RLock()
	peer, ok := g.peers[addr]
	g.mu.RUnlock()
	if !ok {
		return errors.New("can't call RPC on unconnected peer " + string(addr))
	}

	conn, err := peer.open()
	if err != nil {
		return err
	}
	defer conn.Close()

	// write header
	if err := encoding.WriteObject(conn, handlerName(name)); err != nil {
		return err
	}
	// call fn
	return fn(conn)
}

// RPC calls an RPC on the given address. RPC cannot be called on an address
// that the Gateway is not connected to.
func (g *Gateway) RPC(addr modules.NetAddress, name string, fn modules.RPCFunc) error {
	if err := g.threads.Add(); err != nil {
		return err
	}
	defer g.threads.Done()
	return g.managedRPC(addr, name, fn)
}

// RegisterRPC registers an RPCFunc as a handler for a given identifier. To
// call an RPC, use gateway.RPC, supplying the same identifier given to
// RegisterRPC. Identifiers should always use PascalCase. The first 8
// characters of an identifier should be unique, as the identifier used
// internally is truncated to 8 bytes.
func (g *Gateway) RegisterRPC(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; ok {
		build.Critical("RPC already registered: " + name)
	}
	g.handlers[handlerName(name)] = fn
}

// UnregisterRPC unregisters an RPC and removes the corresponding RPCFunc from
// g.handlers. Future calls to the RPC by peers will fail.
func (g *Gateway) UnregisterRPC(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.handlers[handlerName(name)]; !ok {
		build.Critical("RPC not registered: " + name)
	}
	delete(g.handlers, handlerName(name))
}

// RegisterConnectCall registers a name and RPCFunc to be called on a peer
// upon connecting.
func (g *Gateway) RegisterConnectCall(name string, fn modules.RPCFunc) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; ok {
		build.Critical("ConnectCall already registered: " + name)
	}
	g.initRPCs[name] = fn
}

// UnregisterConnectCall unregisters an on-connect call and removes the
// corresponding RPCFunc from g.initRPCs. Future connections to peers will not
// trigger the RPC to be called on them.
func (g *Gateway) UnregisterConnectCall(name string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.initRPCs[name]; !ok {
		build.Critical("ConnectCall not registered: " + name)
	}
	delete(g.initRPCs, name)
}

// threadedListenPeer listens for new streams on a peer connection and serves them via
// threadedHandleConn.
func (g *Gateway) threadedListenPeer(p *peer) {
	// threadedListenPeer registers to the peerTG instead of the primary thread
	// group because peer connections can be lifetime in length, but can also
	// be short-lived. The fact that they can be lifetime means that they can't
	// call threads.Add as they will block calls to threads.Flush. The fact
	// that they can be short-lived means that threads.OnStop is not a good
	// tool for closing out the threads. Instead, they register to peerTG,
	// which is cleanly closed upon gateway shutdown but will not block any
	// calls to threads.Flush()
	if g.peerTG.Add() != nil {
		return
	}
	defer g.peerTG.Done()

	// Spin up a goroutine to listen for a shutdown signal from both the peer
	// and from the gateway. In the event of either, close the muxado session.
	connClosedChan := make(chan struct{})
	peerCloseChan := make(chan struct{})
	go func() {
		// Signal that the muxado session has been successfully closed, and
		// that this goroutine has terminated.
		defer close(connClosedChan)

		// Listen for a stop signal.
		select {
		case <-g.threads.StopChan():
		case <-peerCloseChan:
		}

		// Can't call Disconnect because it could return sync.ErrStopped.
		g.mu.Lock()
		delete(g.peers, p.NetAddress)
		g.mu.Unlock()
		if err := p.sess.Close(); err != nil {
			g.log.Debugf("WARN: error disconnecting from peer %q: %v", p.NetAddress, err)
		}
	}()

	for {
		conn, err := p.accept()
		if err != nil {
			g.log.Debugln("Peer connection closed:", p.NetAddress)
			break
		}

		// it is the handler's responsibility to close the connection
		go g.threadedHandleConn(conn)
	}
	// Signal that the goroutine can shutdown.
	close(peerCloseChan)
	// Wait for confirmation that the goroutine has shut down before returning
	// and releasing the threadgroup registration.
	<-connClosedChan
}

// threadedHandleConn reads header data from a connection, then routes it to the
// appropriate handler for further processing.
func (g *Gateway) threadedHandleConn(conn modules.PeerConn) {
	defer conn.Close()
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	var id rpcID
	if err := encoding.ReadObject(conn, &id, 8); err != nil {
		return
	}
	// call registered handler for this ID
	g.mu.RLock()
	fn, ok := g.handlers[id]
	g.mu.RUnlock()
	if !ok {
		g.log.Debugf("WARN: incoming conn %v requested unknown RPC \"%v\"", conn.RPCAddr(), id)
		return
	}
	g.log.Debugf("INFO: incoming conn %v requested RPC \"%v\"", conn.RPCAddr(), id)

	// call fn
	err := fn(conn)
	// don't log benign errors
	if err == modules.ErrDuplicateTransactionSet || err == modules.ErrBlockKnown {
		err = nil
	}
	if err != nil {
		g.log.Debugf("WARN: incoming RPC \"%v\" from conn %v failed: %v", id, conn.RPCAddr(), err)
	}
}

// Broadcast calls an RPC on all of the specified peers. The calls are run in
// parallel. Broadcasts are restricted to "one-way" RPCs, which simply write an
// object and disconnect. This is why Broadcast takes an interface{} instead of
// an RPCFunc.
func (g *Gateway) Broadcast(name string, obj interface{}, peers []modules.Peer) {
	if g.threads.Add() != nil {
		return
	}
	defer g.threads.Done()

	g.log.Debugf("INFO: broadcasting RPC %q to %v peers", name, len(peers))

	// only encode obj once, instead of using WriteObject
	enc := encoding.Marshal(obj)
	fn := func(conn modules.PeerConn) error {
		return encoding.WritePrefix(conn, enc)
	}

	var wg sync.WaitGroup
	for _, p := range peers {
		wg.Add(1)
		go func(addr modules.NetAddress) {
			defer wg.Done()
			err := g.managedRPC(addr, name, fn)
			if err != nil {
				g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed (attempting again in 10 seconds): %v", name, addr, err)
				// try one more time before giving up
				select {
				case <-time.After(10 * time.Second):
				case <-g.threads.StopChan():
					return
				}
				err := g.managedRPC(addr, name, fn)
				if err != nil {
					g.log.Debugf("WARN: broadcasting RPC %q to peer %q failed twice: %v", name, addr, err)
				}
			}
		}(p.NetAddress)
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>// Copyright 2013, Cong Ding. All rights reserved.
// Use of this source code is governed by a GPLv2
// license that can be found in the LICENSE file.
//
// author: Cong Ding <dinggnu@gmail.com>
package logging

import (
	"runtime"
	"strconv"
	"sync/atomic"
	"time"
)

type field func(*logging) string

var fields = map[string]field{
	"seqid":           (*logging).seqid,
	"name":            (*logging).name,
	"levelno":         (*logging).levelno,
	"levelname":       (*logging).levelname,
	"pathname":        (*logging).pathname,
	"filename":        (*logging).filename,
	"module":          (*logging).module,
	"lineno":          (*logging).lineno,
	"funcName":        (*logging).funcName,
	"created":         (*logging).created,
	"asctime":         (*logging).asctime,
	"msecs":           (*logging).msecs,
	"relativeCreated": (*logging).relativeCreated,
	"thread":          (*logging).thread,
	"threadName":      (*logging).threadName,
	"process":         (*logging).process,
	"message":         (*logging).message,
}

// GetGoId returns the id of goroutine, which is defined in ./get_go_id.c
func GetGoId() int32

func init() {
}

func (logger *logging) seqid() string {
	return strconv.FormatUint(atomic.AddUint64(&(logger.Seqid), 1), 10)
}

func (logger *logging) name() string {
	return logger.Name
}

func (logger *logging) levelno() string {
	return strconv.Itoa(int(logger.Level))
}

func (logger *logging) levelname() string {
	return levelNames[logger.Level]
}

func (logger *logging) pathname() string {
	return ""
}

func (logger *logging) filename() string {
	calldepth := 5
	_, file, _, ok := runtime.Caller(calldepth)
	if !ok {
		file = "???"
	}
	return file
}

func (logger *logging) module() string {
	return ""
}

func (logger *logging) lineno() string {
	calldepth := 5
	_, _, line, ok := runtime.Caller(calldepth)
	if !ok {
		line = 0
	}
	return strconv.Itoa(line)
}

func (logger *logging) funcName() string {
	calldepth := 5
	pc, _, _, ok := runtime.Caller(calldepth)
	if !ok {
		return "???"
	}
	return runtime.FuncForPC(pc).Name()
}

func (logger *logging) created() string {
	return strconv.FormatInt(logger.startTime, 10)
}

func (logger *logging) asctime() string {
	return time.Now().String()
}

func (logger *logging) msecs() string {
	return strconv.Itoa(int(logger.startTime % 1000))
}

func (logger *logging) timestamp() string {
	return strconv.FormatInt(time.Now().UnixNano(), 10)
}

func (logger *logging) relativeCreated() string {
	return strconv.FormatInt(time.Now().UnixNano()-logger.startTime, 10)
}

func (logger *logging) thread() string {
	return strconv.Itoa(int(GetGoId()))
}

func (logger *logging) threadName() string {
	return strconv.Itoa(int(GetGoId()))
}

func (logger *logging) process() string {
	return ""
}

func (logger *logging) message() string {
	return ""
}
<commit_msg>add get pathname<commit_after>// Copyright 2013, Cong Ding. All rights reserved.
// Use of this source code is governed by a GPLv2
// license that can be found in the LICENSE file.
//
// author: Cong Ding <dinggnu@gmail.com>
package logging

import (
	"runtime"
	"strconv"
	"sync/atomic"
	"time"
)

type field func(*logging) string

var fields = map[string]field{
	"seqid":           (*logging).seqid,
	"name":            (*logging).name,
	"levelno":         (*logging).levelno,
	"levelname":       (*logging).levelname,
	"pathname":        (*logging).pathname,
	"filename":        (*logging).filename,
	"module":          (*logging).module,
	"lineno":          (*logging).lineno,
	"funcName":        (*logging).funcName,
	"created":         (*logging).created,
	"asctime":         (*logging).asctime,
	"msecs":           (*logging).msecs,
	"relativeCreated": (*logging).relativeCreated,
	"thread":          (*logging).thread,
	"threadName":      (*logging).threadName,
	"process":         (*logging).process,
	"message":         (*logging).message,
}

// GetGoId returns the id of goroutine, which is defined in ./get_go_id.c
func GetGoId() int32

func init() {
}

func (logger *logging) seqid() string {
	return strconv.FormatUint(atomic.AddUint64(&(logger.Seqid), 1), 10)
}

func (logger *logging) name() string {
	return logger.Name
}

func (logger *logging) levelno() string {
	return strconv.Itoa(int(logger.Level))
}

func (logger *logging) levelname() string {
	return levelNames[logger.Level]
}

func (logger *logging) pathname() string {
	calldepth := 5
	_, file, _, ok := runtime.Caller(calldepth)
	if !ok {
		file = "???"
	}
	return file
}

func (logger *logging) filename() string {
	calldepth := 5
	_, file, _, ok := runtime.Caller(calldepth)
	if !ok {
		file = "???"
	}
	short := file
	for i := len(file) - 1; i > 0; i-- {
		if file[i] == '/' {
			short = file[i+1:]
			break
		}
	}
	file = short
	return file
}

func (logger *logging) module() string {
	return ""
}

func (logger *logging) lineno() string {
	calldepth := 5
	_, _, line, ok := runtime.Caller(calldepth)
	if !ok {
		line = 0
	}
	return strconv.Itoa(line)
}

func (logger *logging) funcName() string {
	calldepth := 5
	pc, _, _, ok := runtime.Caller(calldepth)
	if !ok {
		return "???"
	}
	return runtime.FuncForPC(pc).Name()
}

func (logger *logging) created() string {
	return strconv.FormatInt(logger.startTime, 10)
}

func (logger *logging) asctime() string {
	return time.Now().String()
}

func (logger *logging) msecs() string {
	return strconv.Itoa(int(logger.startTime % 1000))
}

func (logger *logging) timestamp() string {
	return strconv.FormatInt(time.Now().UnixNano(), 10)
}

func (logger *logging) relativeCreated() string {
	return strconv.FormatInt(time.Now().UnixNano()-logger.startTime, 10)
}

func (logger *logging) thread() string {
	return strconv.Itoa(int(GetGoId()))
}

func (logger *logging) threadName() string {
	return strconv.Itoa(int(GetGoId()))
}

func (logger *logging) process() string {
	return ""
}

func (logger *logging) message() string {
	return ""
}
<|endoftext|>


<commit_before>//
// Copyright © 2011-2018 Guy M. Allard
//
// 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 stompngo

import (
	"bufio"
	// "fmt"
	"strings"
)

/*
	Connection handler, one time use during initial connect.

	Handle broker response, react to version incompatabilities, set up session,
	and if necessary initialize heart beats.
*/
func (c *Connection) connectHandler(h Headers) (e error) {
	//fmt.Printf("CHDB01\n")
	c.rdr = bufio.NewReader(c.netconn)
	b, e := c.rdr.ReadBytes(0)
	if e != nil {
		return e
	}
	//fmt.Printf("CHDB02\n")
	f, e := connectResponse(string(b))
	if e != nil {
		return e
	}
	//fmt.Printf("CHDB03\n")
	//
	c.ConnectResponse = &Message{f.Command, f.Headers, f.Body}
	if c.ConnectResponse.Command == ERROR {
		return ECONERR
	}
	//fmt.Printf("CHDB04\n")
	//
	e = c.setProtocolLevel(h, c.ConnectResponse.Headers)
	if e != nil {
		return e
	}
	//fmt.Printf("CHDB05\n")
	//
	if s, ok := c.ConnectResponse.Headers.Contains(HK_SESSION); ok {
		c.session = s
	}

	if c.Protocol() >= SPL_11 {
		e = c.initializeHeartBeats(h)
		if e != nil {
			return e
		}
	}
	//fmt.Printf("CHDB06\n")

	c.connected = true
	c.mets.tfr += 1
	c.mets.tbr += c.ConnectResponse.Size(false)
	return nil
}

/*
	Handle data from the wire after CONNECT is sent. Attempt to create a Frame
	from the wire data.

	Called one time per connection at connection start.
*/
func connectResponse(s string) (*Frame, error) {
	//
	f := new(Frame)
	f.Headers = Headers{}
	f.Body = make([]uint8, 0)

	// Get f.Command
	c := strings.SplitN(s, "\n", 2)
	if len(c) < 2 {
		return nil, EBADFRM
	}
	f.Command = c[0]
	if f.Command != CONNECTED && f.Command != ERROR {
		return f, EUNKFRM
	}

	switch c[1] {
	case "\x00", "\n": // No headers, malformed bodies
		f.Body = []uint8(c[1])
		return f, EBADFRM
	case "\n\x00": // No headers, no body is OK
		return f, nil
	default: // Otherwise continue
	}

	b := strings.SplitN(c[1], "\n\n", 2)
	if len(b) == 1 { // No Headers, b[0] == body
		w := []uint8(b[0])
		f.Body = w[0 : len(w)-1]
		if f.Command == CONNECTED && len(f.Body) > 0 {
			return f, EBDYDATA
		}
		return f, nil
	}

	// Here:
	// b[0] - the headers
	// b[1] - the body

	// Get f.Headers
	for _, l := range strings.Split(b[0], "\n") {
		p := strings.SplitN(l, ":", 2)
		if len(p) < 2 {
			f.Body = []uint8(p[0]) // Bad feedback
			return f, EUNKHDR
		}
		f.Headers = append(f.Headers, p[0], p[1])
	}
	// get f.Body
	w := []uint8(b[1])
	f.Body = w[0 : len(w)-1]
	if f.Command == CONNECTED && len(f.Body) > 0 {
		return f, EBDYDATA
	}

	return f, nil
}

/*
	Check client version, one time use during initial connect.
*/
func (c *Connection) checkClientVersions(h Headers) (e error) {
	w := h.Value(HK_ACCEPT_VERSION)
	if w == "" { // Not present, client wants 1.0
		return nil
	}
	v := strings.SplitN(w, ",", -1) //
	ok := false
	for _, sv := range v {
		if hasValue(supported, sv) {
			ok = true // At least we support one the client wants
		}
	}
	if !ok {
		return EBADVERCLI
	}
	if _, ok = h.Contains(HK_HOST); !ok {
		return EREQHOST
	}
	return nil
}

/*
	Set the protocol level for this new connection.
*/
func (c *Connection) setProtocolLevel(ch, sh Headers) (e error) {
	chw := ch.Value(HK_ACCEPT_VERSION)
	shr := sh.Value(HK_VERSION)

	if chw == shr && Supported(shr) {
		c.protocol = shr
		return nil
	}
	if chw == "" && shr == "" { // Straight up 1.0
		return nil // protocol level defaults to SPL_10
	}
	cv := strings.SplitN(chw, ",", -1) // Client requested versions

	if chw != "" && shr != "" {
		if hasValue(cv, shr) {
			if !Supported(shr) {
				return EBADVERSVR // Client and server agree, but we do not support it
			}
			c.protocol = shr
			return nil
		} else {
			return EBADVERCLI
		}
	}
	if chw != "" && shr == "" { // Client asked for something, server is pure 1.0
		if hasValue(cv, SPL_10) {
			return nil // protocol level defaults to SPL_10
		}
	}

	c.protocol = shr // Could be anything we support
	return nil
}

/*
	Internal function, used only during CONNECT processing.
*/
func hasValue(a []string, w string) bool {
	for _, v := range a {
		if v == w {
			return true
		}
	}
	return false
}
<commit_msg>When error happens in connect pharse, return error with more details of why server refuse to response to connect<commit_after>//
// Copyright © 2011-2018 Guy M. Allard
//
// 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 stompngo

import (
	"bufio"
	// "fmt"
	"strings"
)

type CONNERROR struct {
	err  error
	desc string
}

func (e *CONNERROR) Error() string {
	return e.err.Error() + ":" + e.desc
}

/*
	Connection handler, one time use during initial connect.

	Handle broker response, react to version incompatabilities, set up session,
	and if necessary initialize heart beats.
*/
func (c *Connection) connectHandler(h Headers) (e error) {
	//fmt.Printf("CHDB01\n")
	c.rdr = bufio.NewReader(c.netconn)
	b, e := c.rdr.ReadBytes(0)
	if e != nil {
		return e
	}
	//fmt.Printf("CHDB02\n")
	f, e := connectResponse(string(b))
	if e != nil {
		return e
	}
	//fmt.Printf("CHDB03\n")
	//
	c.ConnectResponse = &Message{f.Command, f.Headers, f.Body}
	if c.ConnectResponse.Command == ERROR {
		return &CONNERROR{ECONERR, string(f.Body)}
	}
	//fmt.Printf("CHDB04\n")
	//
	e = c.setProtocolLevel(h, c.ConnectResponse.Headers)
	if e != nil {
		return e
	}
	//fmt.Printf("CHDB05\n")
	//
	if s, ok := c.ConnectResponse.Headers.Contains(HK_SESSION); ok {
		c.session = s
	}

	if c.Protocol() >= SPL_11 {
		e = c.initializeHeartBeats(h)
		if e != nil {
			return e
		}
	}
	//fmt.Printf("CHDB06\n")

	c.connected = true
	c.mets.tfr += 1
	c.mets.tbr += c.ConnectResponse.Size(false)
	return nil
}

/*
	Handle data from the wire after CONNECT is sent. Attempt to create a Frame
	from the wire data.

	Called one time per connection at connection start.
*/
func connectResponse(s string) (*Frame, error) {
	//
	f := new(Frame)
	f.Headers = Headers{}
	f.Body = make([]uint8, 0)

	// Get f.Command
	c := strings.SplitN(s, "\n", 2)
	if len(c) < 2 {
		return nil, EBADFRM
	}
	f.Command = c[0]
	if f.Command != CONNECTED && f.Command != ERROR {
		return f, EUNKFRM
	}

	switch c[1] {
	case "\x00", "\n": // No headers, malformed bodies
		f.Body = []uint8(c[1])
		return f, EBADFRM
	case "\n\x00": // No headers, no body is OK
		return f, nil
	default: // Otherwise continue
	}

	b := strings.SplitN(c[1], "\n\n", 2)
	if len(b) == 1 { // No Headers, b[0] == body
		w := []uint8(b[0])
		f.Body = w[0 : len(w)-1]
		if f.Command == CONNECTED && len(f.Body) > 0 {
			return f, EBDYDATA
		}
		return f, nil
	}

	// Here:
	// b[0] - the headers
	// b[1] - the body

	// Get f.Headers
	for _, l := range strings.Split(b[0], "\n") {
		p := strings.SplitN(l, ":", 2)
		if len(p) < 2 {
			f.Body = []uint8(p[0]) // Bad feedback
			return f, EUNKHDR
		}
		f.Headers = append(f.Headers, p[0], p[1])
	}
	// get f.Body
	w := []uint8(b[1])
	f.Body = w[0 : len(w)-1]
	if f.Command == CONNECTED && len(f.Body) > 0 {
		return f, EBDYDATA
	}

	return f, nil
}

/*
	Check client version, one time use during initial connect.
*/
func (c *Connection) checkClientVersions(h Headers) (e error) {
	w := h.Value(HK_ACCEPT_VERSION)
	if w == "" { // Not present, client wants 1.0
		return nil
	}
	v := strings.SplitN(w, ",", -1) //
	ok := false
	for _, sv := range v {
		if hasValue(supported, sv) {
			ok = true // At least we support one the client wants
		}
	}
	if !ok {
		return EBADVERCLI
	}
	if _, ok = h.Contains(HK_HOST); !ok {
		return EREQHOST
	}
	return nil
}

/*
	Set the protocol level for this new connection.
*/
func (c *Connection) setProtocolLevel(ch, sh Headers) (e error) {
	chw := ch.Value(HK_ACCEPT_VERSION)
	shr := sh.Value(HK_VERSION)

	if chw == shr && Supported(shr) {
		c.protocol = shr
		return nil
	}
	if chw == "" && shr == "" { // Straight up 1.0
		return nil // protocol level defaults to SPL_10
	}
	cv := strings.SplitN(chw, ",", -1) // Client requested versions

	if chw != "" && shr != "" {
		if hasValue(cv, shr) {
			if !Supported(shr) {
				return EBADVERSVR // Client and server agree, but we do not support it
			}
			c.protocol = shr
			return nil
		} else {
			return EBADVERCLI
		}
	}
	if chw != "" && shr == "" { // Client asked for something, server is pure 1.0
		if hasValue(cv, SPL_10) {
			return nil // protocol level defaults to SPL_10
		}
	}

	c.protocol = shr // Could be anything we support
	return nil
}

/*
	Internal function, used only during CONNECT processing.
*/
func hasValue(a []string, w string) bool {
	for _, v := range a {
		if v == w {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package console

import (
	"bytes"
	"testing"

	"github.com/stretchr/testify/assert"
)

// TestTableGeneration asserts that printing of table table data matches the expected format.
func TestTableGeneration(t *testing.T) {
	tcs := []struct {
		name     string
		table    *Table
		expected string
	}{
		{
			name:  "without_body",
			table: NewTable(Headers("a", "b", "c")),
			expected: `+---+---+---+
| a | b | c |
+---+---+---+
+---+---+---+
`,
		},
		{
			name:  "one_line_body",
			table: NewTable(Headers("a", "b", "c"), body([][]string{{"x", "y", "z"}})),
			expected: `+---+---+---+
| a | b | c |
+---+---+---+
| x | y | z |
+---+---+---+
`,
		},
		{
			name: "with_padding",
			table: NewTable(Headers("NAME", "IP", "MAC"),
				body([][]string{
					{"host1", "192.168.12.23", "AA:BB:CC:DD:EE:FF"},
					{"host2", "192.168.12.102", "99:99:99:99:99:99"},
					{"fritz.box", "192.168.12.1", "11:22:33:44:55:66"},
				})),
			expected: `+-----------+----------------+-------------------+
|   NAME    |       IP       |        MAC        |
+-----------+----------------+-------------------+
| host1     | 192.168.12.23  | AA:BB:CC:DD:EE:FF |
| host2     | 192.168.12.102 | 99:99:99:99:99:99 |
| fritz.box | 192.168.12.1   | 11:22:33:44:55:66 |
+-----------+----------------+-------------------+
`,
		},
		{
			name: "with_alignment",
			table: NewTable(Headers("NAME", "IP", "MAC", "ACT/ONL", "SPEED [MBIT/S]"),
				body([][]string{
					{"host1", "192.168.12.23", "AA:BB:CC:DD:EE:FF", "✔/✘", "104"},
					{"host2", "192.168.12.102", "99:99:99:99:99:99", "✘/✘", "0"},
					{"fritz.box", "192.168.12.1", "11:22:33:44:55:66 ", "✔/✘", "0"},
				})),
			expected: `+-----------+----------------+-------------------+---------+----------------+
|   NAME    |       IP       |        MAC        | ACT/ONL | SPEED [MBIT/S] |
+-----------+----------------+-------------------+---------+----------------+
| host1     | 192.168.12.23  | AA:BB:CC:DD:EE:FF | ✔/✘     |            104 |
| host2     | 192.168.12.102 | 99:99:99:99:99:99 | ✘/✘     |              0 |
| fritz.box | 192.168.12.1   | 11:22:33:44:55:66 | ✔/✘     |              0 |
+-----------+----------------+-------------------+---------+----------------+
`,
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			assertions := assert.New(t)
			assertions.NotNil(tc.table)
			buf := new(bytes.Buffer)
			tc.table.Print(buf)
			assertions.Equal(tc.expected, buf.String())
		})
	}
}

func body(rows [][]string) Option {
	return func(t *Table) {
		for _, r := range rows {
			t.Append(r)
		}
	}
}
<commit_msg>Issue #77: fix test data<commit_after>package console

import (
	"bytes"
	"testing"

	"github.com/stretchr/testify/assert"
)

// TestTableGeneration asserts that printing of table table data matches the expected format.
func TestTableGeneration(t *testing.T) {
	tcs := []struct {
		name     string
		table    *Table
		expected string
	}{
		{
			name:  "without_body",
			table: NewTable(Headers("a", "b", "c")),
			expected: `+---+---+---+
| a | b | c |
+---+---+---+
+---+---+---+
`,
		},
		{
			name:  "one_line_body",
			table: NewTable(Headers("a", "b", "c"), body([][]string{{"x", "y", "z"}})),
			expected: `+---+---+---+
| a | b | c |
+---+---+---+
| x | y | z |
+---+---+---+
`,
		},
		{
			name: "with_padding",
			table: NewTable(Headers("NAME", "IP", "MAC"),
				body([][]string{
					{"host1", "192.168.12.23", "AA:BB:CC:DD:EE:FF"},
					{"host2", "192.168.12.102", "99:99:99:99:99:99"},
					{"fritz.box", "192.168.12.1", "11:22:33:44:55:66"},
				})),
			expected: `+-----------+----------------+-------------------+
|   NAME    |       IP       |        MAC        |
+-----------+----------------+-------------------+
| host1     | 192.168.12.23  | AA:BB:CC:DD:EE:FF |
| host2     | 192.168.12.102 | 99:99:99:99:99:99 |
| fritz.box | 192.168.12.1   | 11:22:33:44:55:66 |
+-----------+----------------+-------------------+
`,
		},
		{
			name: "with_alignment",
			table: NewTable(Headers("NAME", "IP", "MAC", "ACT/ONL", "SPEED [MBIT/S]"),
				body([][]string{
					{"host1", "192.168.12.23", "AA:BB:CC:DD:EE:FF", "✔/✘", "104"},
					{"host2", "192.168.12.102", "99:99:99:99:99:99", "✘/✘", "0"},
					{"fritz.box", "192.168.12.1", "11:22:33:44:55:66", "✔/✘", "0"},
				})),
			expected: `+-----------+----------------+-------------------+---------+----------------+
|   NAME    |       IP       |        MAC        | ACT/ONL | SPEED [MBIT/S] |
+-----------+----------------+-------------------+---------+----------------+
| host1     | 192.168.12.23  | AA:BB:CC:DD:EE:FF | ✔/✘     |            104 |
| host2     | 192.168.12.102 | 99:99:99:99:99:99 | ✘/✘     |              0 |
| fritz.box | 192.168.12.1   | 11:22:33:44:55:66 | ✔/✘     |              0 |
+-----------+----------------+-------------------+---------+----------------+
`,
		},
	}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			assertions := assert.New(t)
			assertions.NotNil(tc.table)
			buf := new(bytes.Buffer)
			tc.table.Print(buf)
			assertions.Equal(tc.expected, buf.String())
		})
	}
}

func body(rows [][]string) Option {
	return func(t *Table) {
		for _, r := range rows {
			t.Append(r)
		}
	}
}
<|endoftext|>


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

// Functions defined in this file should *ONLY* be used for testing.  These
// functions are exported for testing purposes only, and shouldn't be called
// from code that isn't in a test file.

package lxc

import (
	. "launchpad.net/gocheck"
	"launchpad.net/golxc"
	"launchpad.net/juju-core/container/lxc/mock"
)

// SetContainerDir allows tests in other packages to override the
// containerDir.
func SetContainerDir(dir string) (old string) {
	old, containerDir = containerDir, dir
	return
}

// SetLxcContainerDir allows tests in other packages to override the
// lxcContainerDir.
func SetLxcContainerDir(dir string) (old string) {
	old, lxcContainerDir = lxcContainerDir, dir
	return
}

// SetRemovedContainerDir allows tests in other packages to override the
// removedContainerDir.
func SetRemovedContainerDir(dir string) (old string) {
	old, removedContainerDir = removedContainerDir, dir
	return
}

// SetLxcFactory allows tests in other packages to override the lxcObjectFactory
func SetLxcFactory(factory golxc.ContainerFactory) (old golxc.ContainerFactory) {
	logger.Infof("lxcObjectFactory replaced with %v", factory)
	old, lxcObjectFactory = lxcObjectFactory, factory
	return
}

// TestSuite replaces the lxc factory that the broker uses with a mock
// implementation.
type TestSuite struct {
	Factory            mock.ContainerFactory
	oldFactory         golxc.ContainerFactory
	ContainerDir       string
	RemovedDir         string
	LxcDir             string
	oldContainerDir    string
	oldRemovedDir      string
	oldLxcContainerDir string
}

func (s *TestSuite) SetUpSuite(c *C) {}

func (s *TestSuite) TearDownSuite(c *C) {}

func (s *TestSuite) SetUpTest(c *C) {
	s.ContainerDir = c.MkDir()
	s.oldContainerDir = SetContainerDir(s.ContainerDir)
	s.RemovedDir = c.MkDir()
	s.oldRemovedDir = SetRemovedContainerDir(s.RemovedDir)
	s.LxcDir = c.MkDir()
	s.oldLxcContainerDir = SetLxcContainerDir(s.LxcDir)
	s.Factory = mock.MockFactory()
	s.oldFactory = SetLxcFactory(s.Factory)
}

func (s *TestSuite) TearDownTest(c *C) {
	SetContainerDir(s.oldContainerDir)
	SetLxcContainerDir(s.oldLxcContainerDir)
	SetRemovedContainerDir(s.oldRemovedDir)
	SetLxcFactory(s.oldFactory)
}
<commit_msg>And to the test.go file.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// Functions defined in this file should *ONLY* be used for testing.  These
// functions are exported for testing purposes only, and shouldn't be called
// from code that isn't in a test file.

package lxc

import (
	gc "launchpad.net/gocheck"
	"launchpad.net/golxc"

	"launchpad.net/juju-core/container/lxc/mock"
)

// SetContainerDir allows tests in other packages to override the
// containerDir.
func SetContainerDir(dir string) (old string) {
	old, containerDir = containerDir, dir
	return
}

// SetLxcContainerDir allows tests in other packages to override the
// lxcContainerDir.
func SetLxcContainerDir(dir string) (old string) {
	old, lxcContainerDir = lxcContainerDir, dir
	return
}

// SetRemovedContainerDir allows tests in other packages to override the
// removedContainerDir.
func SetRemovedContainerDir(dir string) (old string) {
	old, removedContainerDir = removedContainerDir, dir
	return
}

// SetLxcFactory allows tests in other packages to override the lxcObjectFactory
func SetLxcFactory(factory golxc.ContainerFactory) (old golxc.ContainerFactory) {
	logger.Infof("lxcObjectFactory replaced with %v", factory)
	old, lxcObjectFactory = lxcObjectFactory, factory
	return
}

// TestSuite replaces the lxc factory that the broker uses with a mock
// implementation.
type TestSuite struct {
	Factory            mock.ContainerFactory
	oldFactory         golxc.ContainerFactory
	ContainerDir       string
	RemovedDir         string
	LxcDir             string
	oldContainerDir    string
	oldRemovedDir      string
	oldLxcContainerDir string
}

func (s *TestSuite) SetUpSuite(c *gc.C) {}

func (s *TestSuite) TearDownSuite(c *gc.C) {}

func (s *TestSuite) SetUpTest(c *gc.C) {
	s.ContainerDir = c.MkDir()
	s.oldContainerDir = SetContainerDir(s.ContainerDir)
	s.RemovedDir = c.MkDir()
	s.oldRemovedDir = SetRemovedContainerDir(s.RemovedDir)
	s.LxcDir = c.MkDir()
	s.oldLxcContainerDir = SetLxcContainerDir(s.LxcDir)
	s.Factory = mock.MockFactory()
	s.oldFactory = SetLxcFactory(s.Factory)
}

func (s *TestSuite) TearDownTest(c *gc.C) {
	SetContainerDir(s.oldContainerDir)
	SetLxcContainerDir(s.oldLxcContainerDir)
	SetRemovedContainerDir(s.oldRemovedDir)
	SetLxcFactory(s.oldFactory)
}
<|endoftext|>


<commit_before>package goprocessctx

import (
	goprocess "github.com/jbenet/goprocess"
	context "golang.org/x/net/context"
)

// WithContext constructs and returns a Process that respects
// given context. It is the equivalent of:
//
//   func ProcessWithContext(ctx context.Context) goprocess.Process {
//     p := goprocess.WithParent(goprocess.Background())
//     go func() {
//       <-ctx.Done()
//       p.Close()
//     }()
//     return p
//   }
//
func WithContext(ctx context.Context) goprocess.Process {
	if ctx == nil {
		panic("nil Context")
	}

	p := goprocess.WithParent(goprocess.Background())
	go func() {
		<-ctx.Done()
		p.Close()
	}()
	return p
}

// WithContextAndTeardown is a helper function to set teardown at initiation
// of WithContext
func WithContextAndTeardown(ctx context.Context, tf goprocess.TeardownFunc) goprocess.Process {
	if ctx == nil {
		panic("nil Context")
	}
	p := goprocess.WithTeardown(tf)
	go func() {
		<-ctx.Done()
		p.Close()
	}()
	return p
}

// WaitForContext makes p WaitFor ctx. When Closing, p waits for
// ctx.Done(), before being Closed(). It is simply:
//
//   p.WaitFor(goprocess.WithContext(ctx))
//
func WaitForContext(ctx context.Context, p goprocess.Process) {
	p.WaitFor(WithContext(ctx))
}

// WithProcessClosing returns a context.Context derived from ctx that
// is cancelled as p is Closing (after: <-p.Closing()). It is simply:
//
//   func WithProcessClosing(ctx context.Context, p goprocess.Process) context.Context {
//     ctx, cancel := context.WithCancel(ctx)
//     go func() {
//       <-p.Closing()
//       cancel()
//     }()
//     return ctx
//   }
//
func WithProcessClosing(ctx context.Context, p goprocess.Process) context.Context {
	ctx, cancel := context.WithCancel(ctx)
	go func() {
		<-p.Closing()
		cancel()
	}()
	return ctx
}

// WithProcessClosed returns a context.Context that is cancelled
// after Process p is Closed. It is the equivalent of:
//
//   func WithProcessClosed(ctx context.Context, p goprocess.Process) context.Context {
//     ctx, cancel := context.WithCancel(ctx)
//     go func() {
//       <-p.Closed()
//       cancel()
//     }()
//     return ctx
//   }
//
func WithProcessClosed(ctx context.Context, p goprocess.Process) context.Context {
	ctx, cancel := context.WithCancel(ctx)
	go func() {
		<-p.Closed()
		cancel()
	}()
	return ctx
}
<commit_msg>ctx: added CloseAfterContext<commit_after>package goprocessctx

import (
	goprocess "github.com/jbenet/goprocess"
	context "golang.org/x/net/context"
)

// WithContext constructs and returns a Process that respects
// given context. It is the equivalent of:
//
//   func ProcessWithContext(ctx context.Context) goprocess.Process {
//     p := goprocess.WithParent(goprocess.Background())
//     CloseAfterContext(p, ctx)
//     return p
//   }
//
func WithContext(ctx context.Context) goprocess.Process {
	p := goprocess.WithParent(goprocess.Background())
	CloseAfterContext(p, ctx)
	return p
}

// WithContextAndTeardown is a helper function to set teardown at initiation
// of WithContext
func WithContextAndTeardown(ctx context.Context, tf goprocess.TeardownFunc) goprocess.Process {
	if ctx == nil {
		panic("nil Context")
	}
	p := goprocess.WithTeardown(tf)
	go func() {
		<-ctx.Done()
		p.Close()
	}()
	return p
}

// WaitForContext makes p WaitFor ctx. When Closing, p waits for
// ctx.Done(), before being Closed(). It is simply:
//
//   p.WaitFor(goprocess.WithContext(ctx))
//
func WaitForContext(ctx context.Context, p goprocess.Process) {
	p.WaitFor(WithContext(ctx))
}

// CloseAfterContext schedules the process to close after the given
// context is done. It is the equivalent of:
//
//   func CloseAfterContext(p goprocess.Process, ctx context.Context) {
//     go func() {
//       <-ctx.Done()
//       p.Close()
//     }()
//   }
//
func CloseAfterContext(p goprocess.Process, ctx context.Context) {
	if p == nil {
		panic("nil Process")
	}
	if ctx == nil {
		panic("nil Context")
	}

	go func() {
		<-ctx.Done()
		p.Close()
	}()
}

// WithProcessClosing returns a context.Context derived from ctx that
// is cancelled as p is Closing (after: <-p.Closing()). It is simply:
//
//   func WithProcessClosing(ctx context.Context, p goprocess.Process) context.Context {
//     ctx, cancel := context.WithCancel(ctx)
//     go func() {
//       <-p.Closing()
//       cancel()
//     }()
//     return ctx
//   }
//
func WithProcessClosing(ctx context.Context, p goprocess.Process) context.Context {
	ctx, cancel := context.WithCancel(ctx)
	go func() {
		<-p.Closing()
		cancel()
	}()
	return ctx
}

// WithProcessClosed returns a context.Context that is cancelled
// after Process p is Closed. It is the equivalent of:
//
//   func WithProcessClosed(ctx context.Context, p goprocess.Process) context.Context {
//     ctx, cancel := context.WithCancel(ctx)
//     go func() {
//       <-p.Closed()
//       cancel()
//     }()
//     return ctx
//   }
//
func WithProcessClosed(ctx context.Context, p goprocess.Process) context.Context {
	ctx, cancel := context.WithCancel(ctx)
	go func() {
		<-p.Closed()
		cancel()
	}()
	return ctx
}
<|endoftext|>


<commit_before>package context

import (
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"strings"

	"github.com/gulien/orbit/helpers"

	"github.com/joho/godotenv"
	"gopkg.in/yaml.v2"
)

type (
	// OrbitContext struct contains the data necessary for generating a file from a template.
	OrbitContext struct {
		// TemplateFilePath is the path of the template.
		TemplateFilePath string
		// Values map contains data from YAML files.
		Values map[string]interface{}
		// EnvFiles map contains pairs from .env files.
		EnvFiles map[string]map[string]string
		// Env map contains pairs from environment variables.
		Env map[string]string
		// Os is the OS name at runtime.
		Os string
	}

	// OrbitFileMap struct represents a parameter given to some flags of an Orbit command.
	// Flags: -v --values, -e --env
	OrbitFileMap struct {
		// Name is the given name of the file.
		Name string
		// Path is the path of the file.
		Path string
	}
)

// NewOrbitContext function instantiates a new OrbitContext.
func NewOrbitContext(templateFilePath string, valuesFiles string, envFiles string) (*OrbitContext, error) {
	// as the template is mandatory, we must check its validity.
	if templateFilePath == "" || helpers.FileDoesNotExist(templateFilePath) {
		return nil, fmt.Errorf("Template file %s does not exist", templateFilePath)
	}

	// let's instantiates our OrbitContext!
	ctx := &OrbitContext{
		TemplateFilePath: templateFilePath,
		Os:               runtime.GOOS,
	}

	// checks if a file with values has been specified.
	if valuesFiles != "" {
		data, err := getValuesMap(valuesFiles)
		if err != nil {
			return nil, err
		}

		ctx.Values = data
	}

	// checks if a .env file has been specified.
	if envFiles != "" {
		data, err := getEnvFilesMap(envFiles)
		if err != nil {
			return nil, err
		}

		ctx.EnvFiles = data
	}

	// last but not least, populates the Env map.
	ctx.Env = getEnvMap()

	return ctx, nil
}

// getValuesMap function retrieves values from YAML files.
func getValuesMap(valuesFiles string) (map[string]interface{}, error) {
	filesMap, err := getFilesMap(valuesFiles)
	if err != nil {
		return nil, err
	}

	valuesMap := make(map[string]interface{})
	for _, f := range filesMap {
		// first, checks if the file exists
		if helpers.FileDoesNotExist(f.Path) {
			return nil, fmt.Errorf("Values file %s does not exist", f.Path)
		}

		// the file containing values must be a valid YAML file.
		if !helpers.IsYAML(f.Path) {
			return nil, fmt.Errorf("Values file %s is not a valid YAML file", f.Path)
		}

		// alright, let's read it to retrieve its data!
		data, err := ioutil.ReadFile(f.Path)
		if err != nil {
			return nil, fmt.Errorf("Failed to read the values file %s:\n%s", f.Path, err)
		}

		// last but not least, parses the YAML.
		var values interface{}
		if err := yaml.Unmarshal(data, &values); err != nil {
			return nil, fmt.Errorf("Values file %s is not a valid YAML file:\n%s", f.Path, err)
		}

		valuesMap[f.Name] = values
	}

	return valuesMap, nil
}

// getEnvFilesMap function retrieves pairs from .env files.
func getEnvFilesMap(envFiles string) (map[string]map[string]string, error) {
	filesMap, err := getFilesMap(envFiles)
	if err != nil {
		return nil, err
	}

	envFilesMap := make(map[string]map[string]string)
	for _, f := range filesMap {
		// first, checks if the file exists
		if helpers.FileDoesNotExist(f.Path) {
			return nil, fmt.Errorf("Env file %s does not exist", f.Path)
		}

		// then parses the .env file to retrieve pairs.
		envFilesMap[f.Name], err = godotenv.Read(f.Path)
		if err != nil {
			return nil, fmt.Errorf("Unable to parse the env file %s:\n%s", f.Path, err)
		}
	}

	return envFilesMap, nil
}

// getEnvMap function retrieves all pairs from environment variables.
func getEnvMap() map[string]string {
	envMap := make(map[string]string)
	for _, e := range os.Environ() {
		pair := strings.Split(e, "=")
		envMap[pair[0]] = pair[1]
	}

	return envMap
}

// getFilesMap function reads a string and populates an array of OrbitFileMap instances.
func getFilesMap(s string) ([]*OrbitFileMap, error) {
	var filesMap []*OrbitFileMap

	// checks if the given string is map of files:
	// if not, considers the string as a path.
	// otherwise tries to populate an array of OrbitFileMap instances.
	parts := strings.Split(s, ";")
	if len(parts) == 1 && len(strings.Split(s, ",")) == 1 {
		filesMap = append(filesMap, &OrbitFileMap{"default", s})
	} else {
		for _, part := range parts {
			data := strings.Split(part, ",")
			if len(data) != 2 {
				return filesMap, fmt.Errorf("Unable to process the files map %s", s)
			}

			filesMap = append(filesMap, &OrbitFileMap{data[0], data[1]})
		}
	}

	return filesMap, nil
}
<commit_msg>typo in comment<commit_after>package context

import (
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"strings"

	"github.com/gulien/orbit/helpers"

	"github.com/joho/godotenv"
	"gopkg.in/yaml.v2"
)

type (
	// OrbitContext struct contains the data necessary for generating a file from a template.
	OrbitContext struct {
		// TemplateFilePath is the path of the template.
		TemplateFilePath string
		// Values map contains data from YAML files.
		Values map[string]interface{}
		// EnvFiles map contains pairs from .env files.
		EnvFiles map[string]map[string]string
		// Env map contains pairs from environment variables.
		Env map[string]string
		// Os is the OS name at runtime.
		Os string
	}

	// OrbitFileMap struct represents a parameter given to some flags of an Orbit command.
	// Flags: -v --values, -e --env
	OrbitFileMap struct {
		// Name is the given name of the file.
		Name string
		// Path is the path of the file.
		Path string
	}
)

// NewOrbitContext function instantiates a new OrbitContext.
func NewOrbitContext(templateFilePath string, valuesFiles string, envFiles string) (*OrbitContext, error) {
	// as the template is mandatory, we must check its validity.
	if templateFilePath == "" || helpers.FileDoesNotExist(templateFilePath) {
		return nil, fmt.Errorf("Template file %s does not exist", templateFilePath)
	}

	// let's instantiates our OrbitContext!
	ctx := &OrbitContext{
		TemplateFilePath: templateFilePath,
		Os:               runtime.GOOS,
	}

	// checks if a file with values has been specified.
	if valuesFiles != "" {
		data, err := getValuesMap(valuesFiles)
		if err != nil {
			return nil, err
		}

		ctx.Values = data
	}

	// checks if a .env file has been specified.
	if envFiles != "" {
		data, err := getEnvFilesMap(envFiles)
		if err != nil {
			return nil, err
		}

		ctx.EnvFiles = data
	}

	// last but not least, populates the Env map.
	ctx.Env = getEnvMap()

	return ctx, nil
}

// getValuesMap function retrieves values from YAML files.
func getValuesMap(valuesFiles string) (map[string]interface{}, error) {
	filesMap, err := getFilesMap(valuesFiles)
	if err != nil {
		return nil, err
	}

	valuesMap := make(map[string]interface{})
	for _, f := range filesMap {
		// first, checks if the file exists
		if helpers.FileDoesNotExist(f.Path) {
			return nil, fmt.Errorf("Values file %s does not exist", f.Path)
		}

		// the file containing values must be a valid YAML file.
		if !helpers.IsYAML(f.Path) {
			return nil, fmt.Errorf("Values file %s is not a valid YAML file", f.Path)
		}

		// alright, let's read it to retrieve its data!
		data, err := ioutil.ReadFile(f.Path)
		if err != nil {
			return nil, fmt.Errorf("Failed to read the values file %s:\n%s", f.Path, err)
		}

		// last but not least, parses the YAML.
		var values interface{}
		if err := yaml.Unmarshal(data, &values); err != nil {
			return nil, fmt.Errorf("Values file %s is not a valid YAML file:\n%s", f.Path, err)
		}

		valuesMap[f.Name] = values
	}

	return valuesMap, nil
}

// getEnvFilesMap function retrieves pairs from .env files.
func getEnvFilesMap(envFiles string) (map[string]map[string]string, error) {
	filesMap, err := getFilesMap(envFiles)
	if err != nil {
		return nil, err
	}

	envFilesMap := make(map[string]map[string]string)
	for _, f := range filesMap {
		// first, checks if the file exists
		if helpers.FileDoesNotExist(f.Path) {
			return nil, fmt.Errorf("Env file %s does not exist", f.Path)
		}

		// then parses the .env file to retrieve pairs.
		envFilesMap[f.Name], err = godotenv.Read(f.Path)
		if err != nil {
			return nil, fmt.Errorf("Unable to parse the env file %s:\n%s", f.Path, err)
		}
	}

	return envFilesMap, nil
}

// getEnvMap function retrieves all pairs from environment variables.
func getEnvMap() map[string]string {
	envMap := make(map[string]string)
	for _, e := range os.Environ() {
		pair := strings.Split(e, "=")
		envMap[pair[0]] = pair[1]
	}

	return envMap
}

// getFilesMap function reads a string and populates an array of OrbitFileMap instances.
func getFilesMap(s string) ([]*OrbitFileMap, error) {
	var filesMap []*OrbitFileMap

	// checks if the given string is a map of files:
	// if not, considers the string as a path.
	// otherwise tries to populate an array of OrbitFileMap instances.
	parts := strings.Split(s, ";")
	if len(parts) == 1 && len(strings.Split(s, ",")) == 1 {
		filesMap = append(filesMap, &OrbitFileMap{"default", s})
	} else {
		for _, part := range parts {
			data := strings.Split(part, ",")
			if len(data) != 2 {
				return filesMap, fmt.Errorf("Unable to process the files map %s", s)
			}

			filesMap = append(filesMap, &OrbitFileMap{data[0], data[1]})
		}
	}

	return filesMap, nil
}
<|endoftext|>


<commit_before>package emails_test

import (
	"github.com/alphagov/paas-cf/tools/user_emails/emails"
	"github.com/alphagov/paas-cf/tools/user_emails/emails/stubs"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

type Csv struct {
	Email string `csv:"email"`
	Org string `csv:"org"`
	Role string `csv:"role"`
}

type userInfo struct {
	Username string
	Role string
}

var _ = Describe("Emails", func() {
	Context("with 'normal' urgency", func(){
		It("gets the spaces for each organisation", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			emails.FetchEmails(&cfFake, false, false)
			Expect(cfFake.ListSpacesByQueryCallCount()).To(Equal(3))
		})

		It("extracts the username of each space developer in each space", func() {
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, false)
			Expect(len(names)).To(Equal(4))
			for _, item := range names {
				Expect(item.Email).ToNot(BeNil())
				Expect(item.Org).ToNot(BeNil())
			}
		})

		It("only returns usernames which are valid email addresses", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, false)
			Expect(names).ToNot(ContainElement("admin"))
		})

		It("catches addresses that are known to be problematic", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, false)
			Expect(len(names)).To(Equal(4))
			for _, item := range names {
				Expect(item.Email).To(ContainElement("test@homeoffice.x.gsi.gov.uk"))
				Expect(item.Org).To(ContainElement("Org 2"))
			}
		})

		//It("de-duplicates email addresses", func(){
		//	_, cfFake := stubs.CreateFakeWithStubData()
		//
		//	names := emails.FetchEmails(&cfFake, false, false)
		//	Expect(names).To(HaveLen(4))
		//
		//	i := 0
		//	for _, value := range names {
		//		if value == "user-1@paas.gov"{
		//			i++
		//		}
		//	}
		//	Expect(i).To(Equal(1))
		//})
	})

	Context("with 'critical' urgency", func(){
		It("includes the username of each org manager", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true,false)
			Expect(names).To(ContainElement("org-1-manager-1@paas.gov"))
			Expect(names).To(ContainElement("org-1-manager-2@paas.gov"))
			Expect(names).To(ContainElement("org-2-manager-1@paas.gov"))
			Expect(names).To(ContainElement("org-3-manager-1@paas.gov"))
		})

		It("includes the username of each org auditor", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true, false)
			Expect(names).To(ContainElement("org-1-auditor-1@paas.gov"))
			Expect(names).To(ContainElement("org-1-auditor-2@paas.gov"))
			Expect(names).To(ContainElement("org-2-auditor-1@paas.gov"))
			Expect(names).To(ContainElement("org-3-auditor-1@paas.gov"))
		})

		It("includes the username of each space manager in each space", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true, false)
			Expect(names).To(ContainElement("org-1-space-1-manager-1@paas.gov"))
			Expect(names).To(ContainElement("org-1-space-1-manager-2@paas.gov"))
			Expect(names).To(ContainElement("org-2-space-1-manager-1@paas.gov"))
			Expect(names).To(ContainElement("org-3-space-1-manager-1@paas.gov"))
		})

		It("includes the username of each space auditor in each space", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true,false)
			Expect(names).To(ContainElement("org-1-space-1-auditor-1@paas.gov"))
			Expect(names).To(ContainElement("org-1-space-1-auditor-2@paas.gov"))
			Expect(names).To(ContainElement("org-2-space-1-auditor-1@paas.gov"))
			Expect(names).To(ContainElement("org-3-space-1-auditor-1@paas.gov"))
		})

		It("includes the username of each space developer in each space", func() {
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true, false)
			Expect(names).To(ContainElement("user-1@paas.gov"))
			Expect(names).To(ContainElement("user-2@paas.gov"))
			Expect(names).To(ContainElement("user-3@paas.gov"))
		})
	})

	Context("with 'management' message", func(){
		It("includes the username of each org manager", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, true)
			Expect(names).To(ContainElement("org-1-manager-1@paas.gov"))
			Expect(names).To(ContainElement("org-1-manager-2@paas.gov"))
			Expect(names).To(ContainElement("org-2-manager-1@paas.gov"))
			Expect(names).To(ContainElement("org-3-manager-1@paas.gov"))
		})

		It("includes the username of each org auditor", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, true)
			Expect(names).To(ContainElement("org-1-auditor-1@paas.gov"))
			Expect(names).To(ContainElement("org-1-auditor-2@paas.gov"))
			Expect(names).To(ContainElement("org-2-auditor-1@paas.gov"))
			Expect(names).To(ContainElement("org-3-auditor-1@paas.gov"))
		})

		It("includes the username of each org auditor", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, true)
			Expect(names).To(ContainElement("org-1-billing-manager-1@paas.gov"))
			Expect(names).To(ContainElement("org-1-billing-manager-2@paas.gov"))
			Expect(names).To(ContainElement("org-2-billing-manager-1@paas.gov"))
			Expect(names).To(ContainElement("org-3-billing-manager-1@paas.gov"))
		})

	})
})
<commit_msg>Add some tests.<commit_after>package emails_test

import (
	"github.com/alphagov/paas-cf/tools/user_emails/emails"
	"github.com/alphagov/paas-cf/tools/user_emails/emails/stubs"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Emails", func() {
	Context("with 'normal' urgency", func(){
		It("gets the spaces for each organisation", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			emails.FetchEmails(&cfFake, false, false, "https://admin.example.com/organisations/", "prod")
			Expect(cfFake.ListSpacesByQueryCallCount()).To(Equal(3))
		})

		It("extracts the username of each space developer in each space", func() {
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, false, "https://admin.example.com/organisations/", "prod")
			Expect(len(names)).To(Equal(5))
			for _, item := range names {
				Expect(item.Email).ToNot(BeNil())
				Expect(item.Org).ToNot(BeNil())
				Expect(item.Role).To(ContainSubstring("Developer"))
				Expect(item.Admin).ToNot(BeNil())
				Expect(item.Region).ToNot(BeNil())
			}
		})

		It("only returns usernames which are valid email addresses", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, false, "https://admin.example.com/organisations/", "prod")
			Expect(names).ToNot(ContainElement("admin"))
		})

		It("catches addresses that are known to be problematic", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, false, "https://admin.example.com/organisations/", "prod")
			Expect(len(names)).To(Equal(5))
			Expect(names[3].Email).To(ContainSubstring("test@homeoffice.x.gsi.gov.uk"))
			Expect(names[3].Org).To(ContainSubstring("Org 2"))
			Expect(names[3].Role).To(ContainSubstring("Developer"))
			Expect(names[3].Admin).To(ContainSubstring("org-2"))
			Expect(names[3].Region).To(ContainSubstring("prod"))
		})
	})

	Context("with 'critical' urgency", func(){
		It("includes the username of each org manager", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true,false, "https://admin.example.com/organisations/", "prod")
			Expect(names[6].Email).To(ContainSubstring("org-1-manager-1@paas.gov"))
			Expect(names[6].Org).To(ContainSubstring("Org 1"))
			Expect(names[6].Role).To(ContainSubstring("Org Manager"))
			Expect(names[6].Admin).To(ContainSubstring("org-1"))
			Expect(names[6].Region).To(ContainSubstring("prod"))
			Expect(names[7].Email).To(ContainSubstring("org-1-manager-2@paas.gov"))
			Expect(names[7].Org).To(ContainSubstring("Org 1"))
			Expect(names[7].Role).To(ContainSubstring("Org Manager"))
			Expect(names[7].Admin).To(ContainSubstring("org-1"))
			Expect(names[7].Region).To(ContainSubstring("prod"))
			Expect(names[14].Email).To(ContainSubstring("org-2-manager-1@paas.gov"))
			Expect(names[14].Org).To(ContainSubstring("Org 2"))
			Expect(names[14].Role).To(ContainSubstring("Org Manager"))
			Expect(names[14].Admin).To(ContainSubstring("org-2"))
			Expect(names[14].Region).To(ContainSubstring("prod"))
			Expect(names[19].Email).To(ContainSubstring("org-3-manager-1@paas.gov"))
			Expect(names[19].Org).To(ContainSubstring("Org 3"))
			Expect(names[19].Role).To(ContainSubstring("Org Manager"))
			Expect(names[19].Admin).To(ContainSubstring("org-3"))
			Expect(names[19].Region).To(ContainSubstring("prod"))
		})

		It("includes the username of each org auditor", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true, false, "https://admin.example.com/organisations/", "prod")
			Expect(names[8].Email).To(ContainSubstring("org-1-auditor-1@paas.gov"))
			Expect(names[8].Org).To(ContainSubstring("Org 1"))
			Expect(names[8].Role).To(ContainSubstring("Org Auditor"))
			Expect(names[8].Admin).To(ContainSubstring("org-1"))
			Expect(names[8].Region).To(ContainSubstring("prod"))
			Expect(names[9].Email).To(ContainSubstring("org-1-auditor-2@paas.gov"))
			Expect(names[9].Org).To(ContainSubstring("Org 1"))
			Expect(names[9].Role).To(ContainSubstring("Org Auditor"))
			Expect(names[9].Admin).To(ContainSubstring("org-1"))
			Expect(names[9].Region).To(ContainSubstring("prod"))
			Expect(names[15].Email).To(ContainSubstring("org-2-auditor-1@paas.gov"))
			Expect(names[15].Org).To(ContainSubstring("Org 2"))
			Expect(names[15].Role).To(ContainSubstring("Org Auditor"))
			Expect(names[15].Admin).To(ContainSubstring("org-2"))
			Expect(names[15].Region).To(ContainSubstring("prod"))
			Expect(names[20].Email).To(ContainSubstring("org-3-auditor-1@paas.gov"))
			Expect(names[20].Org).To(ContainSubstring("Org 3"))
			Expect(names[20].Role).To(ContainSubstring("Org Auditor"))
			Expect(names[20].Admin).To(ContainSubstring("org-3"))
			Expect(names[20].Region).To(ContainSubstring("prod"))
		})

		It("includes the username of each space manager in each space", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true, false, "https://admin.example.com/organisations/", "prod")
			Expect(names[2].Email).To(ContainSubstring("org-1-space-1-manager-1@paas.gov"))
			Expect(names[2].Org).To(ContainSubstring("Org 1"))
			Expect(names[2].Role).To(ContainSubstring("Space Manager"))
			Expect(names[2].Admin).To(ContainSubstring("org-1"))
			Expect(names[2].Region).To(ContainSubstring("prod"))
			Expect(names[3].Email).To(ContainSubstring("org-1-space-1-manager-2@paas.gov"))
			Expect(names[3].Org).To(ContainSubstring("Org 1"))
			Expect(names[3].Role).To(ContainSubstring("Space Manager"))
			Expect(names[3].Admin).To(ContainSubstring("org-1"))
			Expect(names[3].Region).To(ContainSubstring("prod"))
			Expect(names[12].Email).To(ContainSubstring("org-2-space-1-manager-1@paas.gov"))
			Expect(names[12].Org).To(ContainSubstring("Org 2"))
			Expect(names[12].Role).To(ContainSubstring("Space Manager"))
			Expect(names[12].Admin).To(ContainSubstring("org-2"))
			Expect(names[12].Region).To(ContainSubstring("prod"))
			Expect(names[17].Email).To(ContainSubstring("org-3-space-1-manager-1@paas.gov"))
			Expect(names[17].Org).To(ContainSubstring("Org 3"))
			Expect(names[17].Role).To(ContainSubstring("Space Manager"))
			Expect(names[17].Admin).To(ContainSubstring("org-3"))
			Expect(names[17].Region).To(ContainSubstring("prod"))
		})

		It("includes the username of each space auditor in each space", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true,false, "https://admin.example.com/organisations/", "prod")
			Expect(names[4].Email).To(ContainSubstring("org-1-space-1-auditor-1@paas.gov"))
			Expect(names[4].Org).To(ContainSubstring("Org 1"))
			Expect(names[4].Role).To(ContainSubstring("Space Auditor"))
			Expect(names[4].Admin).To(ContainSubstring("org-1"))
			Expect(names[4].Region).To(ContainSubstring("prod"))
			Expect(names[5].Email).To(ContainSubstring("org-1-space-1-auditor-2@paas.gov"))
			Expect(names[5].Org).To(ContainSubstring("Org 1"))
			Expect(names[5].Role).To(ContainSubstring("Space Auditor"))
			Expect(names[5].Admin).To(ContainSubstring("org-1"))
			Expect(names[5].Region).To(ContainSubstring("prod"))
			Expect(names[13].Email).To(ContainSubstring("org-2-space-1-auditor-1@paas.gov"))
			Expect(names[13].Org).To(ContainSubstring("Org 2"))
			Expect(names[13].Role).To(ContainSubstring("Space Auditor"))
			Expect(names[13].Admin).To(ContainSubstring("org-2"))
			Expect(names[13].Region).To(ContainSubstring("prod"))
			Expect(names[18].Email).To(ContainSubstring("org-3-space-1-auditor-1@paas.gov"))
			Expect(names[18].Org).To(ContainSubstring("Org 3"))
			Expect(names[18].Role).To(ContainSubstring("Space Auditor"))
			Expect(names[18].Admin).To(ContainSubstring("org-3"))
			Expect(names[18].Region).To(ContainSubstring("prod"))
		})

		It("includes the username of each space developer in each space", func() {
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, true, false, "https://admin.example.com/organisations/", "prod")
			Expect(names[0].Email).To(ContainSubstring("user-1@paas.gov"))
			Expect(names[0].Org).To(ContainSubstring("Org 1"))
			Expect(names[0].Role).To(ContainSubstring("Developer"))
			Expect(names[0].Admin).To(ContainSubstring("org-1"))
			Expect(names[0].Region).To(ContainSubstring("prod"))
			Expect(names[1].Email).To(ContainSubstring("user-2@paas.gov"))
			Expect(names[1].Org).To(ContainSubstring("Org 1"))
			Expect(names[1].Role).To(ContainSubstring("Developer"))
			Expect(names[1].Admin).To(ContainSubstring("org-1"))
			Expect(names[1].Region).To(ContainSubstring("prod"))
			Expect(names[16].Email).To(ContainSubstring("user-3@paas.gov"))
			Expect(names[16].Org).To(ContainSubstring("Org 3"))
			Expect(names[16].Role).To(ContainSubstring("Developer"))
			Expect(names[16].Admin).To(ContainSubstring("org-3"))
			Expect(names[16].Region).To(ContainSubstring("prod"))
		})
	})

	Context("with 'management' message", func(){
		It("includes the username of each org manager", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, true, "https://admin.example.com/organisations/", "prod")
			Expect(names[0].Email).To(ContainSubstring("org-1-manager-1@paas.gov"))
			Expect(names[0].Org).To(ContainSubstring("Org 1"))
			Expect(names[0].Role).To(ContainSubstring("Org Manager"))
			Expect(names[0].Admin).To(ContainSubstring("org-1"))
			Expect(names[0].Region).To(ContainSubstring("prod"))
			Expect(names[1].Email).To(ContainSubstring("org-1-manager-2@paas.gov"))
			Expect(names[1].Org).To(ContainSubstring("Org 1"))
			Expect(names[1].Role).To(ContainSubstring("Org Manager"))
			Expect(names[1].Admin).To(ContainSubstring("org-1"))
			Expect(names[1].Region).To(ContainSubstring("prod"))
			Expect(names[6].Email).To(ContainSubstring("org-2-manager-1@paas.gov"))
			Expect(names[6].Org).To(ContainSubstring("Org 2"))
			Expect(names[6].Role).To(ContainSubstring("Org Manager"))
			Expect(names[6].Admin).To(ContainSubstring("org-2"))
			Expect(names[6].Region).To(ContainSubstring("prod"))
			Expect(names[9].Email).To(ContainSubstring("org-3-manager-1@paas.gov"))
			Expect(names[9].Org).To(ContainSubstring("Org 3"))
			Expect(names[9].Role).To(ContainSubstring("Org Manager"))
			Expect(names[9].Admin).To(ContainSubstring("org-3"))
			Expect(names[9].Region).To(ContainSubstring("prod"))
		})

		It("includes the username of each org auditor", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, true, "https://admin.example.com/organisations/", "prod")
			Expect(names[2].Email).To(ContainSubstring("org-1-auditor-1@paas.gov"))
			Expect(names[2].Org).To(ContainSubstring("Org 1"))
			Expect(names[2].Role).To(ContainSubstring("Org Auditor"))
			Expect(names[2].Admin).To(ContainSubstring("org-1"))
			Expect(names[2].Region).To(ContainSubstring("prod"))
			Expect(names[3].Email).To(ContainSubstring("org-1-auditor-2@paas.gov"))
			Expect(names[3].Org).To(ContainSubstring("Org 1"))
			Expect(names[3].Role).To(ContainSubstring("Org Auditor"))
			Expect(names[3].Admin).To(ContainSubstring("org-1"))
			Expect(names[3].Region).To(ContainSubstring("prod"))
			Expect(names[7].Email).To(ContainSubstring("org-2-auditor-1@paas.gov"))
			Expect(names[7].Org).To(ContainSubstring("Org 2"))
			Expect(names[7].Role).To(ContainSubstring("Org Auditor"))
			Expect(names[7].Admin).To(ContainSubstring("org-2"))
			Expect(names[7].Region).To(ContainSubstring("prod"))
			Expect(names[10].Email).To(ContainSubstring("org-3-auditor-1@paas.gov"))
			Expect(names[10].Org).To(ContainSubstring("Org 3"))
			Expect(names[10].Role).To(ContainSubstring("Org Auditor"))
			Expect(names[10].Admin).To(ContainSubstring("org-3"))
			Expect(names[10].Region).To(ContainSubstring("prod"))
		})

		It("includes the username of each org auditor", func(){
			_, cfFake := stubs.CreateFakeWithStubData()

			names := emails.FetchEmails(&cfFake, false, true, "https://admin.example.com/organisations/", "prod")
			Expect(names[4].Email).To(ContainSubstring("org-1-billing-manager-1@paas.gov"))
			Expect(names[4].Org).To(ContainSubstring("Org 1"))
			Expect(names[4].Role).To(ContainSubstring("Billing Manager"))
			Expect(names[4].Admin).To(ContainSubstring("org-1"))
			Expect(names[4].Region).To(ContainSubstring("prod"))
			Expect(names[5].Email).To(ContainSubstring("org-1-billing-manager-2@paas.gov"))
			Expect(names[5].Org).To(ContainSubstring("Org 1"))
			Expect(names[5].Role).To(ContainSubstring("Billing Manager"))
			Expect(names[5].Admin).To(ContainSubstring("org-1"))
			Expect(names[5].Region).To(ContainSubstring("prod"))
			Expect(names[8].Email).To(ContainSubstring("org-2-billing-manager-1@paas.gov"))
			Expect(names[8].Org).To(ContainSubstring("Org 2"))
			Expect(names[8].Role).To(ContainSubstring("Billing Manager"))
			Expect(names[8].Admin).To(ContainSubstring("org-2"))
			Expect(names[8].Region).To(ContainSubstring("prod"))
			Expect(names[11].Email).To(ContainSubstring("org-3-billing-manager-1@paas.gov"))
			Expect(names[11].Org).To(ContainSubstring("Org 3"))
			Expect(names[11].Role).To(ContainSubstring("Billing Manager"))
			Expect(names[11].Admin).To(ContainSubstring("org-3"))
			Expect(names[11].Region).To(ContainSubstring("prod"))
		})

	})
})
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/url"
	"os"
	"os/exec"
	"strings"
)

var (
	ReleaseId    = flag.Int("id", 0, "github/hawser Release ID")
	uploadUrlFmt = "https://uploads.github.com/repos/github/hawser/releases/%d/assets?%s"
)

func mainRelease() {
	if *ReleaseId < 1 {
		fmt.Println("Need a valid github/hawser release id.")
		fmt.Println("usage: script/release -id")
		os.Exit(1)
	}

	file, err := os.Open("bin/releases/build_matrix.json")
	if err != nil {
		fmt.Println("Error opening build_matrix.json:", err)
		fmt.Println("Ensure `script/bootstrap -all` has completed successfully")
		os.Exit(1)
	}

	defer file.Close()

	buildMatrix := make(map[string]Release)
	if err := json.NewDecoder(file).Decode(&buildMatrix); err != nil {
		fmt.Println("Error reading build_matrix.json:", err)
		os.Exit(1)
	}

	for _, rel := range buildMatrix {
		release(rel)
		fmt.Println()
	}
}

func release(rel Release) {
	query := url.Values{}
	query.Add("name", rel.Filename)
	query.Add("label", rel.Label)

	args := []string{
		"-in",
		"-H", "Content-Type: application/octet-stream",
		"-X", "POST",
		"--data-binary", "@bin/releases/" + rel.Filename,
		fmt.Sprintf(uploadUrlFmt, *ReleaseId, query.Encode()),
	}

	fmt.Println("curl", strings.Join(args, " "))

	cmd := exec.Command("curl", args...)

	by, err := cmd.Output()
	if err != nil {
		fmt.Println("Error running curl:", err)
		os.Exit(1)
	}

	fmt.Println(string(by))
}
<commit_msg>アアー アアアア アーアア<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/url"
	"os"
	"os/exec"
	"strings"
)

var (
	ReleaseId    = flag.Int("id", 0, "github/hawser Release ID")
	uploadUrlFmt = "https://uploads.github.com/repos/hawser/git-hawser/releases/%d/assets?%s"
)

func mainRelease() {
	if *ReleaseId < 1 {
		fmt.Println("Need a valid github/hawser release id.")
		fmt.Println("usage: script/release -id")
		os.Exit(1)
	}

	file, err := os.Open("bin/releases/build_matrix.json")
	if err != nil {
		fmt.Println("Error opening build_matrix.json:", err)
		fmt.Println("Ensure `script/bootstrap -all` has completed successfully")
		os.Exit(1)
	}

	defer file.Close()

	buildMatrix := make(map[string]Release)
	if err := json.NewDecoder(file).Decode(&buildMatrix); err != nil {
		fmt.Println("Error reading build_matrix.json:", err)
		os.Exit(1)
	}

	for _, rel := range buildMatrix {
		release(rel)
		fmt.Println()
	}
}

func release(rel Release) {
	query := url.Values{}
	query.Add("name", rel.Filename)
	query.Add("label", rel.Label)

	args := []string{
		"-in",
		"-H", "Content-Type: application/octet-stream",
		"-X", "POST",
		"--data-binary", "@bin/releases/" + rel.Filename,
		fmt.Sprintf(uploadUrlFmt, *ReleaseId, query.Encode()),
	}

	fmt.Println("curl", strings.Join(args, " "))

	cmd := exec.Command("curl", args...)

	by, err := cmd.Output()
	if err != nil {
		fmt.Println("Error running curl:", err)
		os.Exit(1)
	}

	fmt.Println(string(by))
}
<|endoftext|>


<commit_before>package teleinfo

import (
	"os"
	"sync"
	"syscall"

	"github.com/Sirupsen/logrus"
	"github.com/yroffin/jarvis-go-ext/server/utils/logger"
)

// On linux
// apply: stty 1200 cs7 evenp cstopb -igncr -inlcr -brkint -icrnl -opost -isig -icanon -iexten -F /dev/ttyUSB0
// to configure teleinfo tty
// Cf. https://hallard.me/gestion-de-la-teleinfo-avec-un-raspberry-pi-et-une-carte-arduipi/

// Teleinfo : instance Teleinfo device struct
type Teleinfo struct {
	entries map[string]string
}

// Trame EDF
type TeleinfoTrame struct {
	etiquette string // ETIQUETTE (4 à 8 caractères)
	data      string // DATA (1 à 12 caractères)
	hecksum   string // CHECKSUM (caractère de contrôle : Somme de caractère)
}

var instance *Teleinfo
var once sync.Once

// GetInstance : singleton instance
func GetInstance() *Teleinfo {
	once.Do(func() {
		instance = new(Teleinfo)
		instance.init()
	})
	return instance
}

// declare canal
var canal = make(chan byte, 5)

// handleReadFile : read file
func handleReadFile(device string) error {

	logger.NewLogger().WithFields(logrus.Fields{
		"device": device,
	}).Info("handleReadFile")

	s, _ := os.OpenFile(device, syscall.O_RDWR|syscall.O_NOCTTY|syscall.O_NONBLOCK, 0666)

	// Receive reply
	for {
		buf := make([]byte, 128)
		var len, err = s.Read(buf)
		if err != nil { // err will equal io.EOF
			break
		}
		for i := 0; i < len; i++ {
			canal <- buf[i]
		}
	}

	logger.NewLogger().WithFields(logrus.Fields{
		"status": "done",
	}).Info("handleReadFile")

	return nil
}

/**
* phase1: LF (0x0A)
* phase2: ETIQUETTE (4 à 8 caractères)
* phase3: SP (0x20)
* phase4: DATA (1 à 12 caractères)
* phase5: SP (0x20)
* phase6: CHECKSUM (caractère de contrôle : Somme de caractère)
* CR (0x0D)
 */

// submit trame
func submit(teleinfo *Teleinfo, trame TeleinfoTrame) {
	teleinfo.entries[trame.etiquette] = trame.data
}

// single trame
func handleTrame(teleinfo *Teleinfo, trame string) {
	var espace int
	var send TeleinfoTrame
	for i := 0; i < len(trame); i++ {
		switch {
		case trame[i] == 0x20:
			espace++
			continue
		default:
			if espace == 0 {
				send.etiquette += string([]byte{trame[i]})
			}
			if espace == 1 {
				send.data += string([]byte{trame[i]})
			}
			continue
		}
	}
	// submit new value
	submit(teleinfo, send)
}

// all trames detection
func handleTrames(teleinfo *Teleinfo, trame string) {
	var send string
	for i := 0; i < len(trame); i++ {
		switch {
		case trame[i] == 0x0A:
			send = ""
			continue
		case trame[i] == 0x0D:
			handleTrame(teleinfo, send)
			continue
		default:
			send += string([]byte{trame[i]})
		}
	}
}

// worker to consume file
func worker(teleinfo *Teleinfo) {
	var trame string
	var etx bool

	// wait for ETX 0x003
	for i := 0; etx == false; i++ {
		var x = <-canal
		if x == 0x03 {
			etx = true
		}
	}

	// daemon
	for {
		var x = <-canal
		switch {
		case x == 0x03:
			// wait for ETX 0x003
			handleTrames(teleinfo, trame)
			break
		case x == 0x02:
			// wait for STX 0x002
			trame = ""
			break
		case x != 0x02 && x != 0x03:
			// other
			trame += string([]byte{x})
			break
		}
	}
}

var mapMutex sync.Once

// get values
func (teleinfo *Teleinfo) GetEntries(entries map[string]string) map[string]string {
	for key, value := range teleinfo.entries {
		entries[key] = value
	}
	return entries
}

// Init : Init
func (teleinfo *Teleinfo) init() {
	// add map
	teleinfo.entries = make(map[string]string)

	// start worker
	go handleReadFile(getTeleinfoFile())
	go worker(teleinfo)
	logger.NewLogger().WithFields(logrus.Fields{}).Info("Init ok")
}
<commit_msg>Error handlng<commit_after>package teleinfo

import (
	"os"
	"sync"
	"syscall"

	"github.com/Sirupsen/logrus"
	"github.com/yroffin/jarvis-go-ext/server/utils/logger"
)

// On linux
// apply: stty 1200 cs7 evenp cstopb -igncr -inlcr -brkint -icrnl -opost -isig -icanon -iexten -F /dev/ttyUSB0
// to configure teleinfo tty
// Cf. https://hallard.me/gestion-de-la-teleinfo-avec-un-raspberry-pi-et-une-carte-arduipi/

// Teleinfo : instance Teleinfo device struct
type Teleinfo struct {
	entries map[string]string
}

// Trame EDF
type TeleinfoTrame struct {
	etiquette string // ETIQUETTE (4 à 8 caractères)
	data      string // DATA (1 à 12 caractères)
	hecksum   string // CHECKSUM (caractère de contrôle : Somme de caractère)
}

var instance *Teleinfo
var once sync.Once

// GetInstance : singleton instance
func GetInstance() *Teleinfo {
	once.Do(func() {
		instance = new(Teleinfo)
		instance.init()
	})
	return instance
}

// declare canal
var canal = make(chan byte, 5)

// handleReadFile : read file
func handleReadFile(device string) error {

	logger.NewLogger().WithFields(logrus.Fields{
		"device": device,
	}).Info("handleReadFile")

	s, err := os.OpenFile(device, syscall.O_RDONLY|syscall.O_NOCTTY, 0666)

	if err != nil {
		logger.NewLogger().WithFields(logrus.Fields{
			"error": err,
		}).Error("handleReadFile")
	}

	// Receive reply
	for {
		buf := make([]byte, 128)
		var len, err = s.Read(buf)
		if err != nil { // err will equal io.EOF
			break
		}
		for i := 0; i < len; i++ {
			canal <- buf[i]
		}
	}

	logger.NewLogger().WithFields(logrus.Fields{
		"status": "done",
	}).Info("handleReadFile")

	return nil
}

/**
* phase1: LF (0x0A)
* phase2: ETIQUETTE (4 à 8 caractères)
* phase3: SP (0x20)
* phase4: DATA (1 à 12 caractères)
* phase5: SP (0x20)
* phase6: CHECKSUM (caractère de contrôle : Somme de caractère)
* CR (0x0D)
 */

// submit trame
func submit(teleinfo *Teleinfo, trame TeleinfoTrame) {
	teleinfo.entries[trame.etiquette] = trame.data
}

// single trame
func handleTrame(teleinfo *Teleinfo, trame string) {
	var espace int
	var send TeleinfoTrame
	for i := 0; i < len(trame); i++ {
		switch {
		case trame[i] == 0x20:
			espace++
			continue
		default:
			if espace == 0 {
				send.etiquette += string([]byte{trame[i]})
			}
			if espace == 1 {
				send.data += string([]byte{trame[i]})
			}
			continue
		}
	}
	// submit new value
	submit(teleinfo, send)
}

// all trames detection
func handleTrames(teleinfo *Teleinfo, trame string) {
	var send string
	for i := 0; i < len(trame); i++ {
		switch {
		case trame[i] == 0x0A:
			send = ""
			continue
		case trame[i] == 0x0D:
			handleTrame(teleinfo, send)
			continue
		default:
			send += string([]byte{trame[i]})
		}
	}
}

// worker to consume file
func worker(teleinfo *Teleinfo) {
	var trame string
	var etx bool

	// wait for ETX 0x003
	for i := 0; etx == false; i++ {
		var x = <-canal
		if x == 0x03 {
			etx = true
		}
	}

	// daemon
	for {
		var x = <-canal
		switch {
		case x == 0x03:
			// wait for ETX 0x003
			handleTrames(teleinfo, trame)
			break
		case x == 0x02:
			// wait for STX 0x002
			trame = ""
			break
		case x != 0x02 && x != 0x03:
			// other
			trame += string([]byte{x})
			break
		}
	}
}

var mapMutex sync.Once

// get values
func (teleinfo *Teleinfo) GetEntries(entries map[string]string) map[string]string {
	for key, value := range teleinfo.entries {
		entries[key] = value
	}
	return entries
}

// Init : Init
func (teleinfo *Teleinfo) init() {
	// add map
	teleinfo.entries = make(map[string]string)

	// start worker
	go handleReadFile(getTeleinfoFile())
	go worker(teleinfo)
	logger.NewLogger().WithFields(logrus.Fields{}).Info("Init ok")
}
<|endoftext|>


<commit_before>package main

// This utility moves a network interface into a network namespace and
// configures it. The configuration is passed in as a JSON object
// (marshalled prot.NetworkAdapter).  It is necessary to implement
// this as a separate utility as in Go one does not have tight control
// over which OS thread a given Go thread/routing executes but as can
// only enter a namespace with a specific OS thread.
//
// Note, this logs to stdout so that the caller (gcs) can log the
// output itself.

import (
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"runtime"

	"github.com/Microsoft/opengcs/service/gcs/prot"
	log "github.com/sirupsen/logrus"
	"github.com/vishvananda/netlink"
	"github.com/vishvananda/netns"
)

func netnsConfigMain() {
	if err := netnsConfig(); err != nil {
		log.Fatal("netnsConfig returned: ", err)
	}
	os.Exit(0)
}

func netnsConfig() error {
	ifStr := flag.String("if", "", "Interface/Adapter to move/configure")
	nspid := flag.Int("nspid", -1, "Process ID (to locate netns")
	cfgStr := flag.String("cfg", "", "Adapter configuration (json)")

	flag.Parse()
	if *ifStr == "" || *nspid == -1 || *cfgStr == "" {
		return fmt.Errorf("All three arguments must be specified")
	}

	var a prot.NetworkAdapter
	if err := json.Unmarshal([]byte(*cfgStr), &a); err != nil {
		return err
	}

	if a.NatEnabled {
		log.Infof("Configure %s in %d with: %s/%d gw=%s", *ifStr, *nspid, a.AllocatedIPAddress, a.HostIPPrefixLength, a.HostIPAddress)
	} else {
		log.Infof("Configure %s in %s with DHCP", *ifStr, *nspid)
	}

	// Lock the OS Thread so we don't accidentally switch namespaces
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	// Stash current network namespace away and make sure we enter it as we leave
	origNS, err := netns.Get()
	if err != nil {
		return fmt.Errorf("netns.Get() failed: %v", err)
	}
	defer origNS.Close()

	// Get a reference to the new network namespace
	ns, err := netns.GetFromPid(*nspid)
	if err != nil {
		return fmt.Errorf("netns.GetFromPid(%d) failed: %v", *nspid, err)
	}
	defer ns.Close()

	// Get a reference to the interface and make sure it's down
	link, err := netlink.LinkByName(*ifStr)
	if err != nil {
		return fmt.Errorf("netlink.LinkByName(%s) failed: %v", *ifStr, err)
	}
	if err := netlink.LinkSetDown(link); err != nil {
		return fmt.Errorf("netlink.LinkSetDown(%#v) failed: %v", link, err)
	}

	// Move the interface to the new network namespace
	if err := netlink.LinkSetNsPid(link, *nspid); err != nil {
		return fmt.Errorf("netlink.SetNsPid(%#v, %d) failed: %v", link, *nspid, err)
	}

	log.Infof("Switching from %v to %v", origNS, ns)

	// Enter the new network namespace
	if err := netns.Set(ns); err != nil {
		return fmt.Errorf("netns.Set() failed: %v", err)
	}

	// Re-Get a reference to the interface (it may be a different ID in the new namespace)
	link, err = netlink.LinkByName(*ifStr)
	if err != nil {
		return fmt.Errorf("netlink.LinkByName(%s) failed: %v", *ifStr, err)
	}

	// User requested non-default MTU size
	if a.EncapOverhead != 0 {
		mtu := link.Attrs().MTU - int(a.EncapOverhead)
		if err = netlink.LinkSetMTU(link, mtu); err != nil {
			return fmt.Errorf("netlink.LinkSetMTU(%#v, %d) failed: %v", link, mtu, err)
		}
	}

	// Configure the interface
	if a.NatEnabled {
		metric := 1
		if a.EnableLowMetric {
			metric = 500
		}

		// Bring the interface up
		if err := netlink.LinkSetUp(link); err != nil {
			return fmt.Errorf("netlink.LinkSetUp(%#v) failed: %v", link, err)
		}
		// Set IP address
		addr := &net.IPNet{
			IP: net.ParseIP(a.AllocatedIPAddress),
			// TODO(rn): This assumes/hardcodes IPv4
			Mask: net.CIDRMask(int(a.HostIPPrefixLength), 32)}
		ipAddr := &netlink.Addr{IPNet: addr, Label: ""}
		if err := netlink.AddrAdd(link, ipAddr); err != nil {
			return fmt.Errorf("netlink.AddrAdd(%#v, %#v) failed: %v", link, ipAddr, err)
		}
		// Set gateway
		if a.HostIPAddress != "" {
			gw := net.ParseIP(a.HostIPAddress)
			route := netlink.Route{
				Scope:     netlink.SCOPE_UNIVERSE,
				LinkIndex: link.Attrs().Index,
				Gw:        gw,
				Priority:  metric, // This is what ip route add does
			}
			if err := netlink.RouteAdd(&route); err != nil {
				return fmt.Errorf("netlink.RouteAdd(%#v) failed: %v", route, err)
			}
		}
	} else {
		err := exec.Command(
			"udhcpc",
			"-q",
			"-i", *ifStr,
			"-s", "/sbin/udhcpc_config.script").Run()
		if err != nil {
			return fmt.Errorf("udhcpc failed: %v", err)
		}
	}

	// Add some debug logging
	curNS, _ := netns.Get()
	// Refresh link attributes/state
	link, _ = netlink.LinkByIndex(link.Attrs().Index)
	attr := link.Attrs()
	addrs, _ := netlink.AddrList(link, 0)
	log.Infof("%v: %s[idx=%d,type=%s] is %v", curNS, attr.Name, attr.Index, link.Type(), attr.OperState)
	for _, addr := range addrs {
		log.Infof("  %v", addr)
	}

	return nil
}
<commit_msg>added support for setting a gateway ip that's outside a subnet<commit_after>package main

// This utility moves a network interface into a network namespace and
// configures it. The configuration is passed in as a JSON object
// (marshalled prot.NetworkAdapter).  It is necessary to implement
// this as a separate utility as in Go one does not have tight control
// over which OS thread a given Go thread/routing executes but as can
// only enter a namespace with a specific OS thread.
//
// Note, this logs to stdout so that the caller (gcs) can log the
// output itself.

import (
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"runtime"

	"github.com/Microsoft/opengcs/service/gcs/prot"
	log "github.com/sirupsen/logrus"
	"github.com/vishvananda/netlink"
	"github.com/vishvananda/netns"
)

func netnsConfigMain() {
	if err := netnsConfig(); err != nil {
		log.Fatal("netnsConfig returned: ", err)
	}
	os.Exit(0)
}

func netnsConfig() error {
	ifStr := flag.String("if", "", "Interface/Adapter to move/configure")
	nspid := flag.Int("nspid", -1, "Process ID (to locate netns")
	cfgStr := flag.String("cfg", "", "Adapter configuration (json)")

	flag.Parse()
	if *ifStr == "" || *nspid == -1 || *cfgStr == "" {
		return fmt.Errorf("All three arguments must be specified")
	}

	var a prot.NetworkAdapter
	if err := json.Unmarshal([]byte(*cfgStr), &a); err != nil {
		return err
	}

	if a.NatEnabled {
		log.Infof("Configure %s in %d with: %s/%d gw=%s", *ifStr, *nspid, a.AllocatedIPAddress, a.HostIPPrefixLength, a.HostIPAddress)
	} else {
		log.Infof("Configure %s in %s with DHCP", *ifStr, *nspid)
	}

	// Lock the OS Thread so we don't accidentally switch namespaces
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	// Stash current network namespace away and make sure we enter it as we leave
	origNS, err := netns.Get()
	if err != nil {
		return fmt.Errorf("netns.Get() failed: %v", err)
	}
	defer origNS.Close()

	// Get a reference to the new network namespace
	ns, err := netns.GetFromPid(*nspid)
	if err != nil {
		return fmt.Errorf("netns.GetFromPid(%d) failed: %v", *nspid, err)
	}
	defer ns.Close()

	// Get a reference to the interface and make sure it's down
	link, err := netlink.LinkByName(*ifStr)
	if err != nil {
		return fmt.Errorf("netlink.LinkByName(%s) failed: %v", *ifStr, err)
	}
	if err := netlink.LinkSetDown(link); err != nil {
		return fmt.Errorf("netlink.LinkSetDown(%#v) failed: %v", link, err)
	}

	// Move the interface to the new network namespace
	if err := netlink.LinkSetNsPid(link, *nspid); err != nil {
		return fmt.Errorf("netlink.SetNsPid(%#v, %d) failed: %v", link, *nspid, err)
	}

	log.Infof("Switching from %v to %v", origNS, ns)

	// Enter the new network namespace
	if err := netns.Set(ns); err != nil {
		return fmt.Errorf("netns.Set() failed: %v", err)
	}

	// Re-Get a reference to the interface (it may be a different ID in the new namespace)
	link, err = netlink.LinkByName(*ifStr)
	if err != nil {
		return fmt.Errorf("netlink.LinkByName(%s) failed: %v", *ifStr, err)
	}

	// User requested non-default MTU size
	if a.EncapOverhead != 0 {
		mtu := link.Attrs().MTU - int(a.EncapOverhead)
		if err = netlink.LinkSetMTU(link, mtu); err != nil {
			return fmt.Errorf("netlink.LinkSetMTU(%#v, %d) failed: %v", link, mtu, err)
		}
	}

	// Configure the interface
	if a.NatEnabled {
		metric := 1
		if a.EnableLowMetric {
			metric = 500
		}

		// Bring the interface up
		if err := netlink.LinkSetUp(link); err != nil {
			return fmt.Errorf("netlink.LinkSetUp(%#v) failed: %v", link, err)
		}
		// Set IP address
		addr := &net.IPNet{
			IP: net.ParseIP(a.AllocatedIPAddress),
			// TODO(rn): This assumes/hardcodes IPv4
			Mask: net.CIDRMask(int(a.HostIPPrefixLength), 32)}
		ipAddr := &netlink.Addr{IPNet: addr, Label: ""}
		if err := netlink.AddrAdd(link, ipAddr); err != nil {
			return fmt.Errorf("netlink.AddrAdd(%#v, %#v) failed: %v", link, ipAddr, err)
		}
		// Set gateway
		if a.HostIPAddress != "" {
			gw := net.ParseIP(a.HostIPAddress)

			if !addr.Contains(gw) {
				// In the case that a gw is not part of the subnet we are setting gw for,
				// a new addr containing this gw address need to be added into the link to avoid getting
				// unreachable error when adding this out-of-subnet gw route
				log.Infof("gw is outside of the subnet: Configure %s in %d with: %s/%d gw=%s\n",
					*ifStr, *nspid, a.AllocatedIPAddress, a.HostIPPrefixLength, a.HostIPAddress)
				addr2 := &net.IPNet{
					IP:   net.ParseIP(a.HostIPAddress),
					Mask: net.CIDRMask(32, 32)} // This assumes/hardcodes IPv4
				ipAddr2 := &netlink.Addr{IPNet: addr2, Label: ""}
				if err := netlink.AddrAdd(link, ipAddr2); err != nil {
					return fmt.Errorf("netlink.AddrAdd(%#v, %#v) failed: %v", link, ipAddr2, err)
				}
			}
			route := netlink.Route{
				Scope:     netlink.SCOPE_UNIVERSE,
				LinkIndex: link.Attrs().Index,
				Gw:        gw,
				Priority:  metric, // This is what ip route add does
			}
			if err := netlink.RouteAdd(&route); err != nil {
				return fmt.Errorf("netlink.RouteAdd(%#v) failed: %v", route, err)
			}
		}
	} else {
		err := exec.Command(
			"udhcpc",
			"-q",
			"-i", *ifStr,
			"-s", "/sbin/udhcpc_config.script").Run()
		if err != nil {
			return fmt.Errorf("udhcpc failed: %v", err)
		}
	}

	// Add some debug logging
	curNS, _ := netns.Get()
	// Refresh link attributes/state
	link, _ = netlink.LinkByIndex(link.Attrs().Index)
	attr := link.Attrs()
	addrs, _ := netlink.AddrList(link, 0)
	log.Infof("%v: %s[idx=%d,type=%s] is %v", curNS, attr.Name, attr.Index, link.Type(), attr.OperState)
	for _, addr := range addrs {
		log.Infof("  %v", addr)
	}

	return nil
}
<|endoftext|>