<commit_before>package session

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"time"

	"github.com/mafredri/cdp"
	"github.com/mafredri/cdp/internal/errors"
	"github.com/mafredri/cdp/protocol/target"
	"github.com/mafredri/cdp/rpcc"
)

// session represents a session connection to a target.
type session struct {
	ID       target.SessionID
	TargetID target.ID
	recvC    chan []byte
	send     func([]byte) error

	init chan struct{} // Protect conn from early read.
	conn *rpcc.Conn
}

// Ensure that session implements rpcc.Codec.
var _ rpcc.Codec = (*session)(nil)

// WriteRequest implements rpcc.Codec.
func (s *session) WriteRequest(r *rpcc.Request) error {
	data, err := json.Marshal(r)
	if err != nil {
		return err
	}
	return s.send(data)
}

// ReadResponse implements rpcc.Codec.
func (s *session) ReadResponse(r *rpcc.Response) error {
	<-s.init

	select {
	case m := <-s.recvC:
		return json.Unmarshal(m, r)
	case <-s.conn.Context().Done():
		return s.conn.Context().Err()
	}
}

// Conn returns the underlying *rpcc.Conn that uses session as codec.
func (s *session) Conn() *rpcc.Conn { return s.conn }

// Write forwards a target message to the session connection.
// When write returns an error, the session is closed.
func (s *session) Write(data []byte) error {
	select {
	case s.recvC <- data:
		return nil
	case <-s.conn.Context().Done():
		return s.conn.Context().Err()
	}
}

// Close closes the underlying *rpcc.Conn.
func (s *session) Close() error {
	return s.conn.Close()
}

var (
	// We only handle Close on conn to detach the session. The codec
	// handles the actual transport (Read / Write) in this case.
	sessionDetachConn = func(detach func() error) rpcc.DialOption {
		return rpcc.WithDialer(
			func(_ context.Context, _ string) (io.ReadWriteCloser, error) {
				return &closeConn{close: detach}, nil
			},
		)
	}
	sessionCodec = func(s *session) rpcc.DialOption {
		return rpcc.WithCodec(func(_ io.ReadWriter) rpcc.Codec {
			return s
		})
	}
)

// dial attaches to the target via the provided *cdp.Client and creates
// a lightweight RPC connection to the target. Communication is done via
// the underlying *rpcc.Conn for the provided *cdp.Client.
func dial(ctx context.Context, id target.ID, tc *cdp.Client, detachTimeout time.Duration) (s *session, err error) {
	args := target.NewAttachToTargetArgs(id)
	reply, err := tc.Target.AttachToTarget(ctx, args)
	if err != nil {
		return nil, err
	}

	s = &session{
		TargetID: id,
		ID:       reply.SessionID,
		recvC:    make(chan []byte, 1),
		init:     make(chan struct{}),
		send: func(data []byte) error {
			<-s.init
			// TODO(maf): Use async invocation.
			return tc.Target.SendMessageToTarget(s.conn.Context(),
				target.NewSendMessageToTargetArgs(string(data)).
					SetSessionID(s.ID))
		},
	}

	detach := func() error {
		ctx, cancel := context.WithTimeout(context.Background(), detachTimeout)
		defer cancel()

		err := tc.Target.DetachFromTarget(ctx,
			target.NewDetachFromTargetArgs().SetSessionID(s.ID))

		err = cdp.ErrorCause(err)
		if err == context.DeadlineExceeded {
			return fmt.Errorf("session: detach timed out for session %s", s.ID)
		}
		return errors.Wrapf(err, "session: detach failed for session %s", s.ID)
	}

	s.conn, err = rpcc.DialContext(ctx, "", sessionDetachConn(detach), sessionCodec(s))
	if err != nil {
		return nil, err
	}
	close(s.init)

	return s, nil
}
<commit_msg>session: Explicitly use non-flat session mode for the time being (#106)<commit_after>package session

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"time"

	"github.com/mafredri/cdp"
	"github.com/mafredri/cdp/internal/errors"
	"github.com/mafredri/cdp/protocol/target"
	"github.com/mafredri/cdp/rpcc"
)

// session represents a session connection to a target.
type session struct {
	ID       target.SessionID
	TargetID target.ID
	recvC    chan []byte
	send     func([]byte) error

	init chan struct{} // Protect conn from early read.
	conn *rpcc.Conn
}

// Ensure that session implements rpcc.Codec.
var _ rpcc.Codec = (*session)(nil)

// WriteRequest implements rpcc.Codec.
func (s *session) WriteRequest(r *rpcc.Request) error {
	data, err := json.Marshal(r)
	if err != nil {
		return err
	}
	return s.send(data)
}

// ReadResponse implements rpcc.Codec.
func (s *session) ReadResponse(r *rpcc.Response) error {
	<-s.init

	select {
	case m := <-s.recvC:
		return json.Unmarshal(m, r)
	case <-s.conn.Context().Done():
		return s.conn.Context().Err()
	}
}

// Conn returns the underlying *rpcc.Conn that uses session as codec.
func (s *session) Conn() *rpcc.Conn { return s.conn }

// Write forwards a target message to the session connection.
// When write returns an error, the session is closed.
func (s *session) Write(data []byte) error {
	select {
	case s.recvC <- data:
		return nil
	case <-s.conn.Context().Done():
		return s.conn.Context().Err()
	}
}

// Close closes the underlying *rpcc.Conn.
func (s *session) Close() error {
	return s.conn.Close()
}

var (
	// We only handle Close on conn to detach the session. The codec
	// handles the actual transport (Read / Write) in this case.
	sessionDetachConn = func(detach func() error) rpcc.DialOption {
		return rpcc.WithDialer(
			func(_ context.Context, _ string) (io.ReadWriteCloser, error) {
				return &closeConn{close: detach}, nil
			},
		)
	}
	sessionCodec = func(s *session) rpcc.DialOption {
		return rpcc.WithCodec(func(_ io.ReadWriter) rpcc.Codec {
			return s
		})
	}
)

// dial attaches to the target via the provided *cdp.Client and creates
// a lightweight RPC connection to the target. Communication is done via
// the underlying *rpcc.Conn for the provided *cdp.Client.
func dial(ctx context.Context, id target.ID, tc *cdp.Client, detachTimeout time.Duration) (s *session, err error) {
	args := target.NewAttachToTargetArgs(id)
	// The default of this flag will change to true, so until CDP
	// uses flat session mode, we set this to false explicitly.
	// See https://bugs.chromium.org/p/chromium/issues/detail?id=991325.
	args.SetFlatten(false)
	reply, err := tc.Target.AttachToTarget(ctx, args)
	if err != nil {
		return nil, err
	}

	s = &session{
		TargetID: id,
		ID:       reply.SessionID,
		recvC:    make(chan []byte, 1),
		init:     make(chan struct{}),
		send: func(data []byte) error {
			<-s.init
			// TODO(maf): Use async invocation.
			return tc.Target.SendMessageToTarget(s.conn.Context(),
				target.NewSendMessageToTargetArgs(string(data)).
					SetSessionID(s.ID))
		},
	}

	detach := func() error {
		ctx, cancel := context.WithTimeout(context.Background(), detachTimeout)
		defer cancel()

		err := tc.Target.DetachFromTarget(ctx,
			target.NewDetachFromTargetArgs().SetSessionID(s.ID))

		err = cdp.ErrorCause(err)
		if err == context.DeadlineExceeded {
			return fmt.Errorf("session: detach timed out for session %s", s.ID)
		}
		return errors.Wrapf(err, "session: detach failed for session %s", s.ID)
	}

	s.conn, err = rpcc.DialContext(ctx, "", sessionDetachConn(detach), sessionCodec(s))
	if err != nil {
		return nil, err
	}
	close(s.init)

	return s, nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Jigsaw Operations LLC
//
// 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
//
//     https://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 shadowsocks

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"time"

	"github.com/Jigsaw-Code/outline-ss-server/metrics"
	onet "github.com/Jigsaw-Code/outline-ss-server/net"

	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

func findAccessKey(clientConn onet.DuplexConn, cipherList map[string]shadowaead.Cipher) (string, onet.DuplexConn, error) {
	if len(cipherList) == 0 {
		return "", nil, errors.New("Empty cipher list")
	}
	// replayBuffer saves the bytes read from shadowConn, in order to allow for replays.
	var replayBuffer bytes.Buffer
	// Try each cipher until we find one that authenticates successfully.
	// This assumes that all ciphers are AEAD.
	// TODO: Reorder list to try previously successful ciphers first for the client IP.
	// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
	for id, cipher := range cipherList {
		logger.Debugf("Trying key %v", id)
		// tmpReader reads first from the replayBuffer and then from clientConn if it needs more
		// bytes. All bytes read from clientConn are saved in replayBuffer for future replays.
		tmpReader := io.MultiReader(bytes.NewReader(replayBuffer.Bytes()), io.TeeReader(clientConn, &replayBuffer))
		cipherReader := NewShadowsocksReader(tmpReader, cipher)
		// Read should read just enough data to authenticate the payload size.
		_, err := cipherReader.Read(make([]byte, 0))
		if err != nil {
			logger.Debugf("Failed key %v: %v", id, err)
			continue
		}
		logger.Debugf("Selected key %v", id)
		// We don't need to keep storing and replaying the bytes anymore, but we don't want to drop
		// those already read into the replayBuffer.
		ssr := NewShadowsocksReader(io.MultiReader(&replayBuffer, clientConn), cipher)
		ssw := NewShadowsocksWriter(clientConn, cipher)
		return id, onet.WrapConn(clientConn, ssr, ssw).(onet.DuplexConn), nil
	}
	return "", nil, fmt.Errorf("could not find valid key")
}

type tcpService struct {
	listener  *net.TCPListener
	ciphers   *map[string]shadowaead.Cipher
	m         metrics.ShadowsocksMetrics
	isRunning bool
}

func NewTCPService(listener *net.TCPListener, ciphers *map[string]shadowaead.Cipher, m metrics.ShadowsocksMetrics) TCPService {
	return &tcpService{listener: listener, ciphers: ciphers, m: m}
}

type TCPService interface {
	Start()
	Stop() error
}

func (s *tcpService) Start() {
	s.isRunning = true
	for s.isRunning {
		var clientConn onet.DuplexConn
		clientConn, err := s.listener.AcceptTCP()
		if err != nil {
			if !s.isRunning {
				return
			}
			logger.Errorf("Failed to accept: %v", err)
		}

		go func() (connError *onet.ConnectionError) {
			clientLocation, err := s.m.GetLocation(clientConn.RemoteAddr())
			if err != nil {
				logger.Warningf("Failed location lookup: %v", err)
			}
			logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientConn.RemoteAddr().String())
			s.m.AddOpenTCPConnection(clientLocation)
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in TCP handler: %v", r)
				}
			}()
			connStart := time.Now()
			clientConn.(*net.TCPConn).SetKeepAlive(true)
			keyID := ""
			var proxyMetrics metrics.ProxyMetrics
			clientConn = metrics.MeasureConn(clientConn, &proxyMetrics.ProxyClient, &proxyMetrics.ClientProxy)
			defer func() {
				connEnd := time.Now()
				connDuration := connEnd.Sub(connStart)
				clientConn.Close()
				status := "OK"
				if connError != nil {
					logger.Debugf("TCP Error: %v: %v", connError.Message, connError.Cause)
					status = connError.Status
				}
				logger.Debugf("Done with status %v, duration %v", status, connDuration)
				s.m.AddClosedTCPConnection(clientLocation, keyID, status, proxyMetrics, connDuration)
			}()

			keyID, clientConn, err := findAccessKey(clientConn, *s.ciphers)
			if err != nil {
				return &onet.ConnectionError{"ERR_CIPHER", "Failed to find a valid cipher", err}
			}

			tgtAddr, err := socks.ReadAddr(clientConn)
			if err != nil {
				return &onet.ConnectionError{"ERR_READ_ADDRESS", "Failed to get target address", err}
			}
			tgtTCPAddr, err := net.ResolveTCPAddr("tcp", tgtAddr.String())
			if err != nil {
				return &onet.ConnectionError{"ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err}
			}
			if !tgtTCPAddr.IP.IsGlobalUnicast() {
				return &onet.ConnectionError{"ERR_ADDRESS_INVALID", fmt.Sprintf("Target address is not global unicast: %v", tgtAddr.String()), err}
			}

			tgtTCPConn, err := net.DialTCP("tcp", nil, tgtTCPAddr)
			if err != nil {
				return &onet.ConnectionError{"ERR_CONNECT", "Failed to connect to target", err}
			}
			defer tgtTCPConn.Close()
			tgtTCPConn.SetKeepAlive(true)
			tgtConn := metrics.MeasureConn(tgtTCPConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)

			// TODO: Disable logging in production. This is sensitive.
			logger.Debugf("proxy %s <-> %s", clientConn.RemoteAddr().String(), tgtConn.RemoteAddr().String())
			_, _, err = onet.Relay(clientConn, tgtConn)
			if err != nil {
				return &onet.ConnectionError{"ERR_RELAY", "Failed to relay traffic", err}
			}
			return nil
		}()
	}
}

func (s *tcpService) Stop() error {
	s.isRunning = false
	return s.listener.Close()
}
<commit_msg>Optimize debug logging<commit_after>// Copyright 2018 Jigsaw Operations LLC
//
// 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
//
//     https://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 shadowsocks

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"time"

	logging "github.com/op/go-logging"

	"github.com/Jigsaw-Code/outline-ss-server/metrics"
	onet "github.com/Jigsaw-Code/outline-ss-server/net"

	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

func findAccessKey(clientConn onet.DuplexConn, cipherList map[string]shadowaead.Cipher) (string, onet.DuplexConn, error) {
	if len(cipherList) == 0 {
		return "", nil, errors.New("Empty cipher list")
	}
	// replayBuffer saves the bytes read from shadowConn, in order to allow for replays.
	var replayBuffer bytes.Buffer
	// Try each cipher until we find one that authenticates successfully.
	// This assumes that all ciphers are AEAD.
	// TODO: Reorder list to try previously successful ciphers first for the client IP.
	// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
	for id, cipher := range cipherList {
		if logger.IsEnabledFor(logging.DEBUG) {
			logger.Debugf("Trying key %v", id)
		}
		// tmpReader reads first from the replayBuffer and then from clientConn if it needs more
		// bytes. All bytes read from clientConn are saved in replayBuffer for future replays.
		tmpReader := io.MultiReader(bytes.NewReader(replayBuffer.Bytes()), io.TeeReader(clientConn, &replayBuffer))
		cipherReader := NewShadowsocksReader(tmpReader, cipher)
		// Read should read just enough data to authenticate the payload size.
		_, err := cipherReader.Read(make([]byte, 0))
		if err != nil {
			if logger.IsEnabledFor(logging.DEBUG) {
				logger.Debugf("Failed key %v: %v", id, err)
			}
			continue
		}
		if logger.IsEnabledFor(logging.DEBUG) {
			logger.Debugf("Selected key %v", id)
		}
		// We don't need to keep storing and replaying the bytes anymore, but we don't want to drop
		// those already read into the replayBuffer.
		ssr := NewShadowsocksReader(io.MultiReader(&replayBuffer, clientConn), cipher)
		ssw := NewShadowsocksWriter(clientConn, cipher)
		return id, onet.WrapConn(clientConn, ssr, ssw).(onet.DuplexConn), nil
	}
	return "", nil, fmt.Errorf("could not find valid key")
}

type tcpService struct {
	listener  *net.TCPListener
	ciphers   *map[string]shadowaead.Cipher
	m         metrics.ShadowsocksMetrics
	isRunning bool
}

func NewTCPService(listener *net.TCPListener, ciphers *map[string]shadowaead.Cipher, m metrics.ShadowsocksMetrics) TCPService {
	return &tcpService{listener: listener, ciphers: ciphers, m: m}
}

type TCPService interface {
	Start()
	Stop() error
}

func (s *tcpService) Start() {
	s.isRunning = true
	for s.isRunning {
		var clientConn onet.DuplexConn
		clientConn, err := s.listener.AcceptTCP()
		if err != nil {
			if !s.isRunning {
				return
			}
			logger.Errorf("Failed to accept: %v", err)
		}

		go func() (connError *onet.ConnectionError) {
			clientLocation, err := s.m.GetLocation(clientConn.RemoteAddr())
			if err != nil {
				logger.Warningf("Failed location lookup: %v", err)
			}
			logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientConn.RemoteAddr().String())
			s.m.AddOpenTCPConnection(clientLocation)
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in TCP handler: %v", r)
				}
			}()
			connStart := time.Now()
			clientConn.(*net.TCPConn).SetKeepAlive(true)
			keyID := ""
			var proxyMetrics metrics.ProxyMetrics
			clientConn = metrics.MeasureConn(clientConn, &proxyMetrics.ProxyClient, &proxyMetrics.ClientProxy)
			defer func() {
				connEnd := time.Now()
				connDuration := connEnd.Sub(connStart)
				clientConn.Close()
				status := "OK"
				if connError != nil {
					logger.Debugf("TCP Error: %v: %v", connError.Message, connError.Cause)
					status = connError.Status
				}
				logger.Debugf("Done with status %v, duration %v", status, connDuration)
				s.m.AddClosedTCPConnection(clientLocation, keyID, status, proxyMetrics, connDuration)
			}()

			keyID, clientConn, err := findAccessKey(clientConn, *s.ciphers)
			if err != nil {
				return &onet.ConnectionError{"ERR_CIPHER", "Failed to find a valid cipher", err}
			}

			tgtAddr, err := socks.ReadAddr(clientConn)
			if err != nil {
				return &onet.ConnectionError{"ERR_READ_ADDRESS", "Failed to get target address", err}
			}
			tgtTCPAddr, err := net.ResolveTCPAddr("tcp", tgtAddr.String())
			if err != nil {
				return &onet.ConnectionError{"ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err}
			}
			if !tgtTCPAddr.IP.IsGlobalUnicast() {
				return &onet.ConnectionError{"ERR_ADDRESS_INVALID", fmt.Sprintf("Target address is not global unicast: %v", tgtAddr.String()), err}
			}

			tgtTCPConn, err := net.DialTCP("tcp", nil, tgtTCPAddr)
			if err != nil {
				return &onet.ConnectionError{"ERR_CONNECT", "Failed to connect to target", err}
			}
			defer tgtTCPConn.Close()
			tgtTCPConn.SetKeepAlive(true)
			tgtConn := metrics.MeasureConn(tgtTCPConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)

			// TODO: Disable logging in production. This is sensitive.
			logger.Debugf("proxy %s <-> %s", clientConn.RemoteAddr().String(), tgtConn.RemoteAddr().String())
			_, _, err = onet.Relay(clientConn, tgtConn)
			if err != nil {
				return &onet.ConnectionError{"ERR_RELAY", "Failed to relay traffic", err}
			}
			return nil
		}()
	}
}

func (s *tcpService) Stop() error {
	s.isRunning = false
	return s.listener.Close()
}
<|endoftext|>


<commit_before>// Copyright 2018 Jigsaw Operations LLC
//
// 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
//
//     https://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 shadowsocks

import (
	"errors"
	"fmt"
	"net"
	"runtime/debug"
	"time"

	"github.com/Jigsaw-Code/outline-ss-server/metrics"
	onet "github.com/Jigsaw-Code/outline-ss-server/net"

	"sync"

	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

const udpBufSize = 64 * 1024

// upack decripts src into dst. It tries each cipher until it finds one that authenticates
// correctly. dst and src must not overlap.
func unpack(dst, src []byte, ciphers map[string]shadowaead.Cipher) ([]byte, string, shadowaead.Cipher, error) {
	for id, cipher := range ciphers {
		logger.Debugf("Trying UDP cipher %v", id)
		buf, err := shadowaead.Unpack(dst, src, cipher)
		if err != nil {
			logger.Debugf("Failed UDP cipher %v: %v", id, err)
			continue
		}
		logger.Debugf("Selected UDP cipher %v", id)
		return buf, id, cipher, nil
	}
	return nil, "", nil, errors.New("could not find valid cipher")
}

// Listen on addr for encrypted packets and basically do UDP NAT.
// We take the ciphers as a pointer because it gets replaced on config updates.
func RunUDPService(natTimeout time.Duration, clientConn net.PacketConn, ciphers *map[string]shadowaead.Cipher, m metrics.ShadowsocksMetrics) {
	defer clientConn.Close()

	nm := newNATmap(natTimeout, m)
	cipherBuf := make([]byte, udpBufSize)
	textBuf := make([]byte, udpBufSize)

	for {
		func() (connError *onet.ConnectionError) {
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in UDP loop: %v", r)
					debug.PrintStack()
				}
			}()
			clientLocation := ""
			keyID := ""
			var clientProxyBytes, proxyTargetBytes int
			defer func() {
				status := "OK"
				if connError != nil {
					logger.Debugf("UDP Error: %v: %v", connError.Message, connError.Cause)
					status = connError.Status
				}
				m.AddUDPPacketFromClient(clientLocation, keyID, status, clientProxyBytes, proxyTargetBytes)
			}()
			clientProxyBytes, clientAddr, err := clientConn.ReadFrom(cipherBuf)
			if err != nil {
				return onet.NewConnectionError("ERR_READ", "Failed to read from client", err)
			}
			clientLocation, locErr := m.GetLocation(clientAddr)
			if locErr != nil {
				logger.Warning("Failed location lookup: %v", locErr)
			}
			logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientAddr.String())
			defer logger.Debugf("UDP done with %v", clientAddr.String())
			logger.Debugf("UDP Request from %v with %v bytes", clientAddr, clientProxyBytes)
			buf, keyID, cipher, err := unpack(textBuf, cipherBuf[:clientProxyBytes], *ciphers)
			if err != nil {
				return onet.NewConnectionError("ERR_CIPHER", "Failed to upack data from client", err)
			}

			tgtAddr := socks.SplitAddr(buf)
			if tgtAddr == nil {
				return onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", nil)
			}

			tgtUDPAddr, err := net.ResolveUDPAddr("udp", tgtAddr.String())
			if err != nil {
				return onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err)
			}
			if !tgtUDPAddr.IP.IsGlobalUnicast() {
				return onet.NewConnectionError("ERR_ADDRESS_INVALID", fmt.Sprintf("Target address is not global unicast: %v", tgtAddr.String()), err)
			}

			payload := buf[len(tgtAddr):]

			targetConn := nm.Get(clientAddr.String())
			if targetConn == nil {
				targetConn, err = net.ListenPacket("udp", "")
				if err != nil {
					return onet.NewConnectionError("ERR_CREATE_SOCKET", "Failed to create UDP socket", err)
				}
				nm.Add(clientAddr, clientConn, cipher, targetConn, clientLocation, keyID)
			}
			logger.Debugf("UDP Nat: client %v <-> proxy exit %v", clientAddr, targetConn.LocalAddr())

			proxyTargetBytes, err = targetConn.WriteTo(payload, tgtUDPAddr) // accept only UDPAddr despite the signature
			if err != nil {
				return onet.NewConnectionError("ERR_WRITE", "Failed to write to target", err)
			}
			return nil
		}()
	}
}

// Packet NAT table
type natmap struct {
	sync.RWMutex
	keyConn map[string]net.PacketConn
	timeout time.Duration
	metrics metrics.ShadowsocksMetrics
}

func newNATmap(timeout time.Duration, sm metrics.ShadowsocksMetrics) *natmap {
	m := &natmap{metrics: sm}
	m.keyConn = make(map[string]net.PacketConn)
	m.timeout = timeout
	return m
}

func (m *natmap) Get(key string) net.PacketConn {
	m.RLock()
	defer m.RUnlock()
	return m.keyConn[key]
}

func (m *natmap) set(key string, pc net.PacketConn) {
	m.Lock()
	defer m.Unlock()

	m.keyConn[key] = pc
}

func (m *natmap) del(key string) net.PacketConn {
	m.Lock()
	defer m.Unlock()

	pc, ok := m.keyConn[key]
	if ok {
		delete(m.keyConn, key)
		return pc
	}
	return nil
}

func (m *natmap) Add(clientAddr net.Addr, clientConn net.PacketConn, cipher shadowaead.Cipher, targetConn net.PacketConn, clientLocation, keyID string) {
	m.set(clientAddr.String(), targetConn)

	m.metrics.AddUdpNatEntry()
	go func() {
		timedCopy(clientAddr, clientConn, cipher, targetConn, m.timeout, clientLocation, keyID, m.metrics)
		m.metrics.RemoveUdpNatEntry()
		if pc := m.del(clientAddr.String()); pc != nil {
			pc.Close()
		}
	}()
}

// copy from src to dst at target with read timeout
func timedCopy(clientAddr net.Addr, clientConn net.PacketConn, cipher shadowaead.Cipher, targetConn net.PacketConn,
	timeout time.Duration, clientLocation, keyID string, sm metrics.ShadowsocksMetrics) {
	textBuf := make([]byte, udpBufSize)
	cipherBuf := make([]byte, udpBufSize)

	expired := false
	for !expired {
		var targetProxyBytes, proxyClientBytes int
		connError := func() (connError *onet.ConnectionError) {
			var (
				raddr net.Addr
				err   error
			)
			targetConn.SetReadDeadline(time.Now().Add(timeout))
			targetProxyBytes, raddr, err = targetConn.ReadFrom(textBuf)
			if err != nil {
				if netErr, ok := err.(net.Error); ok {
					if netErr.Timeout() {
						expired = true
						return nil
					}
				}
				return onet.NewConnectionError("ERR_READ", "Failed to read from target", err)
			}

			srcAddr := socks.ParseAddr(raddr.String())
			logger.Debugf("UDP response from %v to %v", srcAddr, clientAddr)
			// Shift data buffer to prepend with srcAddr.
			copy(textBuf[len(srcAddr):], textBuf[:targetProxyBytes])
			copy(textBuf, srcAddr)

			buf, err := shadowaead.Pack(cipherBuf, textBuf[:len(srcAddr)+targetProxyBytes], cipher)
			if err != nil {
				return onet.NewConnectionError("ERR_PACK", "Failed to pack data to client", err)
			}
			proxyClientBytes, err = clientConn.WriteTo(buf, clientAddr)
			if err != nil {
				return onet.NewConnectionError("ERR_WRITE", "Failed to write to client", err)
			}
			return nil
		}()
		status := "OK"
		if connError != nil {
			logger.Debugf("UDP Error: %v: %v", connError.Message, connError.Cause)
			status = connError.Status
		}
		sm.AddUDPPacketFromTarget(clientLocation, keyID, status, targetProxyBytes, proxyClientBytes)
	}
}
<commit_msg>Fix UDP warning<commit_after>// Copyright 2018 Jigsaw Operations LLC
//
// 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
//
//     https://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 shadowsocks

import (
	"errors"
	"fmt"
	"net"
	"runtime/debug"
	"time"

	"github.com/Jigsaw-Code/outline-ss-server/metrics"
	onet "github.com/Jigsaw-Code/outline-ss-server/net"

	"sync"

	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

const udpBufSize = 64 * 1024

// upack decripts src into dst. It tries each cipher until it finds one that authenticates
// correctly. dst and src must not overlap.
func unpack(dst, src []byte, ciphers map[string]shadowaead.Cipher) ([]byte, string, shadowaead.Cipher, error) {
	for id, cipher := range ciphers {
		logger.Debugf("Trying UDP cipher %v", id)
		buf, err := shadowaead.Unpack(dst, src, cipher)
		if err != nil {
			logger.Debugf("Failed UDP cipher %v: %v", id, err)
			continue
		}
		logger.Debugf("Selected UDP cipher %v", id)
		return buf, id, cipher, nil
	}
	return nil, "", nil, errors.New("could not find valid cipher")
}

// Listen on addr for encrypted packets and basically do UDP NAT.
// We take the ciphers as a pointer because it gets replaced on config updates.
func RunUDPService(natTimeout time.Duration, clientConn net.PacketConn, ciphers *map[string]shadowaead.Cipher, m metrics.ShadowsocksMetrics) {
	defer clientConn.Close()

	nm := newNATmap(natTimeout, m)
	cipherBuf := make([]byte, udpBufSize)
	textBuf := make([]byte, udpBufSize)

	for {
		func() (connError *onet.ConnectionError) {
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in UDP loop: %v", r)
					debug.PrintStack()
				}
			}()
			clientLocation := ""
			keyID := ""
			var clientProxyBytes, proxyTargetBytes int
			defer func() {
				status := "OK"
				if connError != nil {
					logger.Debugf("UDP Error: %v: %v", connError.Message, connError.Cause)
					status = connError.Status
				}
				m.AddUDPPacketFromClient(clientLocation, keyID, status, clientProxyBytes, proxyTargetBytes)
			}()
			clientProxyBytes, clientAddr, err := clientConn.ReadFrom(cipherBuf)
			if err != nil {
				return onet.NewConnectionError("ERR_READ", "Failed to read from client", err)
			}
			clientLocation, locErr := m.GetLocation(clientAddr)
			if locErr != nil {
				logger.Warningf("Failed location lookup: %v", locErr)
			}
			logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientAddr.String())
			defer logger.Debugf("UDP done with %v", clientAddr.String())
			logger.Debugf("UDP Request from %v with %v bytes", clientAddr, clientProxyBytes)
			buf, keyID, cipher, err := unpack(textBuf, cipherBuf[:clientProxyBytes], *ciphers)
			if err != nil {
				return onet.NewConnectionError("ERR_CIPHER", "Failed to upack data from client", err)
			}

			tgtAddr := socks.SplitAddr(buf)
			if tgtAddr == nil {
				return onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", nil)
			}

			tgtUDPAddr, err := net.ResolveUDPAddr("udp", tgtAddr.String())
			if err != nil {
				return onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err)
			}
			if !tgtUDPAddr.IP.IsGlobalUnicast() {
				return onet.NewConnectionError("ERR_ADDRESS_INVALID", fmt.Sprintf("Target address is not global unicast: %v", tgtAddr.String()), err)
			}

			payload := buf[len(tgtAddr):]

			targetConn := nm.Get(clientAddr.String())
			if targetConn == nil {
				targetConn, err = net.ListenPacket("udp", "")
				if err != nil {
					return onet.NewConnectionError("ERR_CREATE_SOCKET", "Failed to create UDP socket", err)
				}
				nm.Add(clientAddr, clientConn, cipher, targetConn, clientLocation, keyID)
			}
			logger.Debugf("UDP Nat: client %v <-> proxy exit %v", clientAddr, targetConn.LocalAddr())

			proxyTargetBytes, err = targetConn.WriteTo(payload, tgtUDPAddr) // accept only UDPAddr despite the signature
			if err != nil {
				return onet.NewConnectionError("ERR_WRITE", "Failed to write to target", err)
			}
			return nil
		}()
	}
}

// Packet NAT table
type natmap struct {
	sync.RWMutex
	keyConn map[string]net.PacketConn
	timeout time.Duration
	metrics metrics.ShadowsocksMetrics
}

func newNATmap(timeout time.Duration, sm metrics.ShadowsocksMetrics) *natmap {
	m := &natmap{metrics: sm}
	m.keyConn = make(map[string]net.PacketConn)
	m.timeout = timeout
	return m
}

func (m *natmap) Get(key string) net.PacketConn {
	m.RLock()
	defer m.RUnlock()
	return m.keyConn[key]
}

func (m *natmap) set(key string, pc net.PacketConn) {
	m.Lock()
	defer m.Unlock()

	m.keyConn[key] = pc
}

func (m *natmap) del(key string) net.PacketConn {
	m.Lock()
	defer m.Unlock()

	pc, ok := m.keyConn[key]
	if ok {
		delete(m.keyConn, key)
		return pc
	}
	return nil
}

func (m *natmap) Add(clientAddr net.Addr, clientConn net.PacketConn, cipher shadowaead.Cipher, targetConn net.PacketConn, clientLocation, keyID string) {
	m.set(clientAddr.String(), targetConn)

	m.metrics.AddUdpNatEntry()
	go func() {
		timedCopy(clientAddr, clientConn, cipher, targetConn, m.timeout, clientLocation, keyID, m.metrics)
		m.metrics.RemoveUdpNatEntry()
		if pc := m.del(clientAddr.String()); pc != nil {
			pc.Close()
		}
	}()
}

// copy from src to dst at target with read timeout
func timedCopy(clientAddr net.Addr, clientConn net.PacketConn, cipher shadowaead.Cipher, targetConn net.PacketConn,
	timeout time.Duration, clientLocation, keyID string, sm metrics.ShadowsocksMetrics) {
	textBuf := make([]byte, udpBufSize)
	cipherBuf := make([]byte, udpBufSize)

	expired := false
	for !expired {
		var targetProxyBytes, proxyClientBytes int
		connError := func() (connError *onet.ConnectionError) {
			var (
				raddr net.Addr
				err   error
			)
			targetConn.SetReadDeadline(time.Now().Add(timeout))
			targetProxyBytes, raddr, err = targetConn.ReadFrom(textBuf)
			if err != nil {
				if netErr, ok := err.(net.Error); ok {
					if netErr.Timeout() {
						expired = true
						return nil
					}
				}
				return onet.NewConnectionError("ERR_READ", "Failed to read from target", err)
			}

			srcAddr := socks.ParseAddr(raddr.String())
			logger.Debugf("UDP response from %v to %v", srcAddr, clientAddr)
			// Shift data buffer to prepend with srcAddr.
			copy(textBuf[len(srcAddr):], textBuf[:targetProxyBytes])
			copy(textBuf, srcAddr)

			buf, err := shadowaead.Pack(cipherBuf, textBuf[:len(srcAddr)+targetProxyBytes], cipher)
			if err != nil {
				return onet.NewConnectionError("ERR_PACK", "Failed to pack data to client", err)
			}
			proxyClientBytes, err = clientConn.WriteTo(buf, clientAddr)
			if err != nil {
				return onet.NewConnectionError("ERR_WRITE", "Failed to write to client", err)
			}
			return nil
		}()
		status := "OK"
		if connError != nil {
			logger.Debugf("UDP Error: %v: %v", connError.Message, connError.Cause)
			status = connError.Status
		}
		sm.AddUDPPacketFromTarget(clientLocation, keyID, status, targetProxyBytes, proxyClientBytes)
	}
}
<|endoftext|>


<commit_before>package backend

import (
	"errors"

	"context"

	"github.com/joiggama/money"
)

var (
	// ErrNotAnItem is what happens if you try to vend an item what not be there
	ErrNotAnItem = errors.New("Item does not exist")
	// ErrItemEmpty is if you vend something that's out
	ErrItemEmpty = errors.New("Item has no Quantity")
)

// StockItem represents an item in the vending machine
type StockItem struct {
	ID       string
	Name     string
	Quantity uint8
	Reserved uint8
	Image    string
	Price    uint64
	Location uint8
}

// CanVend checks stock availability
func (item *StockItem) CanVend() bool {
	return item.Quantity != 0 && item.Quantity >= item.Reserved
}

var mFormatOptions = money.Options{"currency": "GBP"}

// FormattedPrice returns the price as a currency string
func (item *StockItem) FormattedPrice() string {
	return money.Format(float64(item.Price)/100, mFormatOptions)

}

// Stock is an interface to the vending machine (possibly could be named better)
type Stock struct {
	items map[string]*StockItem
}

func newStock() *Stock {
	return &Stock{
		items: make(map[string]*StockItem, 14),
	}
}

// GetFakeStock returns a Stock with a bunch of fake items in it
func GetFakeStock() *Stock {
	stock := newStock()

	stock.items["c9b2d770-532f-41fc-acf7-e6e682bd82d5"] = &StockItem{
		ID:       "c9b2d770-532f-41fc-acf7-e6e682bd82d5",
		Name:     "Example Item #1",
		Quantity: 5,
		Price:    1000,
		Location: 0,
		Image:    "lorem-pixel-01.jpg",
	}
	stock.items["67fd2453-a2a9-4260-949d-c0f573b4b3ab"] = &StockItem{
		ID:       "67fd2453-a2a9-4260-949d-c0f573b4b3ab",
		Name:     "Example Item #2",
		Quantity: 50,
		Price:    333,
		Location: 1,
		Image:    "lorem-pixel-02.jpg",
	}
	stock.items["47db0c90-e9ff-4b2f-a646-ef259887a89b"] = &StockItem{
		ID:       "47db0c90-e9ff-4b2f-a646-ef259887a89b",
		Name:     "Example Item #3",
		Quantity: 200,
		Price:    11133,
		Location: 2,
		Image:    "lorem-pixel-03.jpg",
	}
	stock.items["aaac9bd8-ac3f-4348-af1c-c211396e3ff9"] = &StockItem{
		ID:       "aaac9bd8-ac3f-4348-af1c-c211396e3ff9",
		Name:     "Example Item #4",
		Quantity: 5,
		Price:    50,
		Location: 3,
		Image:    "lorem-pixel-04.jpg",
	}
	stock.items["552be5d3-ed40-4a2e-bb88-f0319b5f4af1"] = &StockItem{
		ID:       "552be5d3-ed40-4a2e-bb88-f0319b5f4af1",
		Name:     "Example Item #5",
		Quantity: 5,
		Price:    100,
		Location: 4,
		Image:    "lorem-pixel-05.jpg",
	}
	stock.items["ef8e1fc5-8a5d-49fb-baba-37377beead4a"] = &StockItem{
		ID:       "ef8e1fc5-8a5d-49fb-baba-37377beead4a",
		Name:     "Example Item #6",
		Quantity: 5,
		Price:    1200,
		Location: 5,
		Image:    "lorem-pixel-06.jpg",
	}
	stock.items["0840a34d-e024-4e47-a8db-471ac97c6aae"] = &StockItem{
		ID:       "0840a34d-e024-4e47-a8db-471ac97c6aae",
		Name:     "Example Item #7",
		Quantity: 5,
		Price:    123456,
		Location: 6,
		Image:    "lorem-pixel-07.jpg",
	}
	stock.items["193599ed-d1eb-411c-9ecc-e2343256609b"] = &StockItem{
		ID:       "193599ed-d1eb-411c-9ecc-e2343256609b",
		Name:     "Example Item #8",
		Quantity: 5,
		Price:    0,
		Location: 7,
		Image:    "lorem-pixel-08.jpg",
	}
	stock.items["c0e26c01-8f37-4a23-a7e9-caf322036ba9"] = &StockItem{
		ID:       "c0e26c01-8f37-4a23-a7e9-caf322036ba9",
		Name:     "Example Item #9",
		Quantity: 5,
		Price:    100,
		Location: 8,
		Image:    "lorem-pixel-09.jpg",
	}
	stock.items["b4d6ec03-268c-49d6-ac17-c107b3375014"] = &StockItem{
		ID:       "b4d6ec03-268c-49d6-ac17-c107b3375014",
		Name:     "Example Item #10",
		Quantity: 0,
		Price:    1000,
		Location: 9,
		Image:    "lorem-pixel-10.jpg",
	}
	stock.items["31d1fb41-6af1-4fff-ad2e-1138632a6dab"] = &StockItem{
		ID:       "31d1fb41-6af1-4fff-ad2e-1138632a6dab",
		Name:     "Example Item #11",
		Quantity: 5,
		Price:    1,
		Location: 10,
		Image:    "lorem-pixel-11.jpg",
	}
	stock.items["1bb541b6-b465-445e-8bb3-58688bf40e13"] = &StockItem{
		ID:       "1bb541b6-b465-445e-8bb3-58688bf40e13",
		Name:     "Example Item #12",
		Quantity: 8,
		Price:    100,
		Location: 11,
		Image:    "lorem-pixel-12.jpg",
	}
	return stock
}

// GetAll returns all items currently in the vending machine (available or not)
func (stock *Stock) GetAll(ctx context.Context) (items []*StockItem, err error) {
	items = make([]*StockItem, len(stock.items))
	i := 0
	for _, item := range stock.items {
		items[i] = item
		i++
	}

	return
}

// GetItem returns information specific to a single item
func (stock *Stock) GetItem(ctx context.Context, ID string) (item *StockItem, err error) {
	return stock.items[ID], nil
}
<commit_msg>Talk about quantity and reservation atomically<commit_after>package backend

import (
	"errors"
	"sync/atomic"

	"context"

	"github.com/joiggama/money"
)

var (
	// ErrNotAnItem is what happens if you try to vend an item what not be there
	ErrNotAnItem = errors.New("Item does not exist")
	// ErrItemEmpty is if you vend something that's out
	ErrItemEmpty = errors.New("Item has no Quantity")
)

// StockItem represents an item in the vending machine
type StockItem struct {
	ID       string
	Name     string
	Quantity uint32
	Reserved uint32
	Image    string
	Price    uint64
	Location uint8
}

// CanVend checks stock availability
func (item *StockItem) CanVend() bool {
	q := atomic.LoadUint32(&item.Quantity)
	r := atomic.LoadUint32(&item.Reserved)
	// We can only vend if there are unreserved items available
	return q > 0 && q > r
}

var mFormatOptions = money.Options{"currency": "GBP"}

// FormattedPrice returns the price as a currency string
func (item *StockItem) FormattedPrice() string {
	return money.Format(float64(item.Price)/100, mFormatOptions)

}

// Stock is an interface to the vending machine (possibly could be named better)
type Stock struct {
	items map[string]*StockItem
}

func newStock() *Stock {
	return &Stock{
		items: make(map[string]*StockItem, 14),
	}
}

// GetFakeStock returns a Stock with a bunch of fake items in it
func GetFakeStock() *Stock {
	stock := newStock()

	stock.items["c9b2d770-532f-41fc-acf7-e6e682bd82d5"] = &StockItem{
		ID:       "c9b2d770-532f-41fc-acf7-e6e682bd82d5",
		Name:     "Example Item #1",
		Quantity: 5,
		Price:    1000,
		Location: 0,
		Image:    "lorem-pixel-01.jpg",
	}
	stock.items["67fd2453-a2a9-4260-949d-c0f573b4b3ab"] = &StockItem{
		ID:       "67fd2453-a2a9-4260-949d-c0f573b4b3ab",
		Name:     "Example Item #2",
		Quantity: 50,
		Price:    333,
		Location: 1,
		Image:    "lorem-pixel-02.jpg",
	}
	stock.items["47db0c90-e9ff-4b2f-a646-ef259887a89b"] = &StockItem{
		ID:       "47db0c90-e9ff-4b2f-a646-ef259887a89b",
		Name:     "Example Item #3",
		Quantity: 200,
		Price:    11133,
		Location: 2,
		Image:    "lorem-pixel-03.jpg",
	}
	stock.items["aaac9bd8-ac3f-4348-af1c-c211396e3ff9"] = &StockItem{
		ID:       "aaac9bd8-ac3f-4348-af1c-c211396e3ff9",
		Name:     "Example Item #4",
		Quantity: 5,
		Price:    50,
		Location: 3,
		Image:    "lorem-pixel-04.jpg",
	}
	stock.items["552be5d3-ed40-4a2e-bb88-f0319b5f4af1"] = &StockItem{
		ID:       "552be5d3-ed40-4a2e-bb88-f0319b5f4af1",
		Name:     "Example Item #5",
		Quantity: 5,
		Price:    100,
		Location: 4,
		Image:    "lorem-pixel-05.jpg",
	}
	stock.items["ef8e1fc5-8a5d-49fb-baba-37377beead4a"] = &StockItem{
		ID:       "ef8e1fc5-8a5d-49fb-baba-37377beead4a",
		Name:     "Example Item #6",
		Quantity: 5,
		Price:    1200,
		Location: 5,
		Image:    "lorem-pixel-06.jpg",
	}
	stock.items["0840a34d-e024-4e47-a8db-471ac97c6aae"] = &StockItem{
		ID:       "0840a34d-e024-4e47-a8db-471ac97c6aae",
		Name:     "Example Item #7",
		Quantity: 5,
		Price:    123456,
		Location: 6,
		Image:    "lorem-pixel-07.jpg",
	}
	stock.items["193599ed-d1eb-411c-9ecc-e2343256609b"] = &StockItem{
		ID:       "193599ed-d1eb-411c-9ecc-e2343256609b",
		Name:     "Example Item #8",
		Quantity: 5,
		Price:    0,
		Location: 7,
		Image:    "lorem-pixel-08.jpg",
	}
	stock.items["c0e26c01-8f37-4a23-a7e9-caf322036ba9"] = &StockItem{
		ID:       "c0e26c01-8f37-4a23-a7e9-caf322036ba9",
		Name:     "Example Item #9",
		Quantity: 5,
		Price:    100,
		Location: 8,
		Image:    "lorem-pixel-09.jpg",
	}
	stock.items["b4d6ec03-268c-49d6-ac17-c107b3375014"] = &StockItem{
		ID:       "b4d6ec03-268c-49d6-ac17-c107b3375014",
		Name:     "Example Item #10",
		Quantity: 0,
		Price:    1000,
		Location: 9,
		Image:    "lorem-pixel-10.jpg",
	}
	stock.items["31d1fb41-6af1-4fff-ad2e-1138632a6dab"] = &StockItem{
		ID:       "31d1fb41-6af1-4fff-ad2e-1138632a6dab",
		Name:     "Example Item #11",
		Quantity: 5,
		Price:    1,
		Location: 10,
		Image:    "lorem-pixel-11.jpg",
	}
	stock.items["1bb541b6-b465-445e-8bb3-58688bf40e13"] = &StockItem{
		ID:       "1bb541b6-b465-445e-8bb3-58688bf40e13",
		Name:     "Example Item #12",
		Quantity: 8,
		Price:    100,
		Location: 11,
		Image:    "lorem-pixel-12.jpg",
	}
	return stock
}

// GetAll returns all items currently in the vending machine (available or not)
func (stock *Stock) GetAll(ctx context.Context) (items []*StockItem, err error) {
	items = make([]*StockItem, len(stock.items))
	i := 0
	for _, item := range stock.items {
		items[i] = item
		i++
	}

	return
}

// GetItem returns information specific to a single item
func (stock *Stock) GetItem(ctx context.Context, ID string) (item *StockItem, err error) {
	return stock.items[ID], nil
}
<|endoftext|>


<commit_before>package lib

import (
	"bytes"
	"fmt"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/bytemark-client/util/log"
)

//CreateVirtualMachine creates a virtual machine in the given group.
func (c *bytemarkClient) CreateVirtualMachine(group GroupName, spec brain.VirtualMachineSpec) (vm *brain.VirtualMachine, err error) {
	err = c.validateGroupName(&group)
	if err != nil {
		return nil, err
	}
	r, err := c.BuildRequest("POST", BrainEndpoint, "/accounts/%s/groups/%s/vm_create", group.Account, group.Group)
	if err != nil {
		return nil, err
	}
	if spec.IPs != nil {
		if spec.IPs.IPv4 == "" && spec.IPs.IPv6 == "" {
			spec.IPs = nil
		}
	}
	if spec.Discs != nil {
		if len(spec.Discs) == 0 {
			spec.Discs = nil
		}
		for i, disc := range spec.Discs {
			newDisc, discErr := disc.Validate()
			if discErr != nil {
				return nil, discErr
			}
			spec.Discs[i] = *newDisc
		}
		labelDiscs(spec.Discs, 0)
	}

	vm = new(brain.VirtualMachine)
	oldfile := log.LogFile
	log.LogFile = nil
	_, _, err = r.MarshalAndRun(spec, vm)
	log.LogFile = oldfile
	return vm, err
}

// DeleteVirtualMachine deletes the named virtual machine.
// returns nil on success or an error otherwise.
func (c *bytemarkClient) DeleteVirtualMachine(name VirtualMachineName, purge bool) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	purgePart := ""
	if purge {
		purgePart = "?purge=true"
	}
	r, err := c.BuildRequest("DELETE", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s"+purgePart, name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(nil, nil)
	return err
}

// GetVirtualMachine requests an overview of the named VM, regardless of its deletion status.
func (c *bytemarkClient) GetVirtualMachine(name VirtualMachineName) (vm *brain.VirtualMachine, err error) {
	var r *Request

	// If the VM name is numeric, it means it is an internal Bytemark ID,
	// so we should use a different endpoint
	if _, nErr := strconv.Atoi(name.VirtualMachine); nErr == nil {
		r, err = c.BuildRequest("GET", BrainEndpoint, "/virtual_machines/%s?include_deleted=true&view=overview", name.VirtualMachine)
	} else {
		err = c.validateVirtualMachineName(&name)
		if err != nil {
			return
		}
		vm = new(brain.VirtualMachine)
		r, err = c.BuildRequest("GET", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s?include_deleted=true&view=overview", name.Account, name.Group, name.VirtualMachine)
	}
	if err != nil {
		return
	}

	vm = new(brain.VirtualMachine)
	_, _, err = r.Run(nil, vm)
	if err != nil {
		return
	}
	return
}

//MoveVirtualMachine moves the virtual machine to the given name, across groups if needed.
func (c *bytemarkClient) MoveVirtualMachine(oldName VirtualMachineName, newName VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&oldName)
	if err != nil {
		return
	}
	err = c.validateVirtualMachineName(&newName)
	if err != nil {
		return
	}

	// create the change we want to see in the server
	change := brain.VirtualMachine{Name: newName.VirtualMachine}
	if newName.Group != "" || newName.Account != "" {
		// get group
		groupName := GroupName{Group: newName.Group, Account: newName.Account}
		group, groupErr := c.GetGroup(groupName)
		if groupErr != nil {
			return groupErr
		}
		change.GroupID = group.ID
	}

	// PUT the change
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", oldName.Account, oldName.Group, oldName.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.MarshalAndRun(change, nil)
	return err

}

// ReimageVirtualMachine reimages the named virtual machine. This will wipe everything on the first disk in the vm and install a new OS on top of it.
// Note that the machine in question must already be powered off. Once complete, according to the API docs, the vm will be powered on but its autoreboot_on will be false.
func (c *bytemarkClient) ReimageVirtualMachine(name VirtualMachineName, image brain.ImageInstall) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("POST", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s/reimage", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	oldfile := log.LogFile
	log.LogFile = nil
	_, _, err = r.MarshalAndRun(image, nil)
	log.LogFile = oldfile
	return err
}

// ResetVirtualMachine resets the named virtual machine. This is like pressing the reset
// button on a physical computer. This does not cause a new process to be started, so does not apply any pending hardware changes.
// returns nil on success or an error otherwise.
func (c *bytemarkClient) ResetVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("POST", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s/signal", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"signal":"reset"}`), nil)
	return err
}

// RestartVirtualMachine restarts the named virtual machine. This is
// returns nil on success or an error otherwise.
func (c *bytemarkClient) RestartVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"autoreboot_on":true, "power_on": false}`), nil)
	return err
}

// StartVirtualMachine starts the named virtual machine.
// returns nil on success or an error otherwise.
func (c *bytemarkClient) StartVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"autoreboot_on":true, "power_on": true}`), nil)
	return err
}

// StopVirtualMachine starts the named virtual machine.
// returns nil on success or an error otherwise.
func (c *bytemarkClient) StopVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"autoreboot_on":false, "power_on": false}`), nil)
	return err
}

// ShutdownVirtualMachine sends an ACPI shutdown to the VM. This will cause a graceful shutdown of the machine
// returns nil on success or an error otherwise.
func (c *bytemarkClient) ShutdownVirtualMachine(name VirtualMachineName, stayoff bool) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return
	}
	var r *Request
	if stayoff {
		r, err = c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
		if err != nil {
			return
		}

		_, _, err = r.Run(bytes.NewBufferString(`{"autoreboot_on":false}`), nil)
		if err != nil {
			return
		}
	}
	r, err = c.BuildRequest("POST", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s/signal", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"signal": "powerdown"}`), nil)
	return err
}

// UndeleteVirtualMachine changes the deleted flag on a VM back to false.
// Return nil on success, an error otherwise.
func (c *bytemarkClient) UndeleteVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"deleted":false}`), nil)
	return err
}

// SetVirtualMachineHardwareProfile specifies the hardware profile on a VM. Optionally locks or unlocks h. profile
// Return nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineHardwareProfile(name VirtualMachineName, profile string, locked ...bool) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}
	hwprofileLock := ""
	if len(locked) > 0 {
		hwprofileLock = `, "hardware_profile_locked": false`
		if locked[0] {
			hwprofileLock = `, "hardware_profile_locked": true`
		}
	}
	profileJSON := fmt.Sprintf(`{"hardware_profile": "%s"%s}`, profile, hwprofileLock)

	_, _, err = r.Run(bytes.NewBufferString(profileJSON), nil)
	return err
}

// SetVirtualMachineHardwareProfileLock locks or unlocks the hardware profile of a VM.
// Return nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineHardwareProfileLock(name VirtualMachineName, locked bool) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	lockJSON := `{"hardware_profile_locked": false}`
	if locked {
		lockJSON = `{"hardware_profile_locked": true}`
	}

	_, _, err = r.Run(bytes.NewBufferString(lockJSON), nil)
	return err
}

// SetVirtualMachineMemory sets the RAM available to a virtual machine in megabytes
// Return nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineMemory(name VirtualMachineName, memory int) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	memoryJSON := fmt.Sprintf(`{"memory": %d}`, memory)

	_, _, err = r.Run(bytes.NewBufferString(memoryJSON), nil)
	return err
}

// SetVirtualMachineCores sets the number of CPUs available to a virtual machine
// Return nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineCores(name VirtualMachineName, cores int) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	coresJSON := fmt.Sprintf(`{"cores": %d}`, cores)

	_, _, err = r.Run(bytes.NewBufferString(coresJSON), nil)
	return err
}

// SetVirtualMachineCDROM sets the URL of a CD to attach to a virtual machine. Set url to "" to remove the CD.
// Returns nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineCDROM(name VirtualMachineName, url string) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.MarshalAndRun(brain.VirtualMachine{CdromURL: url}, nil)
	return err
}
<commit_msg>Add strconv back in to vm_calls<commit_after>package lib

import (
	"bytes"
	"fmt"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/bytemark-client/util/log"
	"strconv"
)

//CreateVirtualMachine creates a virtual machine in the given group.
func (c *bytemarkClient) CreateVirtualMachine(group GroupName, spec brain.VirtualMachineSpec) (vm *brain.VirtualMachine, err error) {
	err = c.validateGroupName(&group)
	if err != nil {
		return nil, err
	}
	r, err := c.BuildRequest("POST", BrainEndpoint, "/accounts/%s/groups/%s/vm_create", group.Account, group.Group)
	if err != nil {
		return nil, err
	}
	if spec.IPs != nil {
		if spec.IPs.IPv4 == "" && spec.IPs.IPv6 == "" {
			spec.IPs = nil
		}
	}
	if spec.Discs != nil {
		if len(spec.Discs) == 0 {
			spec.Discs = nil
		}
		for i, disc := range spec.Discs {
			newDisc, discErr := disc.Validate()
			if discErr != nil {
				return nil, discErr
			}
			spec.Discs[i] = *newDisc
		}
		labelDiscs(spec.Discs, 0)
	}

	vm = new(brain.VirtualMachine)
	oldfile := log.LogFile
	log.LogFile = nil
	_, _, err = r.MarshalAndRun(spec, vm)
	log.LogFile = oldfile
	return vm, err
}

// DeleteVirtualMachine deletes the named virtual machine.
// returns nil on success or an error otherwise.
func (c *bytemarkClient) DeleteVirtualMachine(name VirtualMachineName, purge bool) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	purgePart := ""
	if purge {
		purgePart = "?purge=true"
	}
	r, err := c.BuildRequest("DELETE", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s"+purgePart, name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(nil, nil)
	return err
}

// GetVirtualMachine requests an overview of the named VM, regardless of its deletion status.
func (c *bytemarkClient) GetVirtualMachine(name VirtualMachineName) (vm *brain.VirtualMachine, err error) {
	var r *Request

	// If the VM name is numeric, it means it is an internal Bytemark ID,
	// so we should use a different endpoint
	if _, nErr := strconv.Atoi(name.VirtualMachine); nErr == nil {
		r, err = c.BuildRequest("GET", BrainEndpoint, "/virtual_machines/%s?include_deleted=true&view=overview", name.VirtualMachine)
	} else {
		err = c.validateVirtualMachineName(&name)
		if err != nil {
			return
		}
		vm = new(brain.VirtualMachine)
		r, err = c.BuildRequest("GET", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s?include_deleted=true&view=overview", name.Account, name.Group, name.VirtualMachine)
	}
	if err != nil {
		return
	}

	vm = new(brain.VirtualMachine)
	_, _, err = r.Run(nil, vm)
	if err != nil {
		return
	}
	return
}

//MoveVirtualMachine moves the virtual machine to the given name, across groups if needed.
func (c *bytemarkClient) MoveVirtualMachine(oldName VirtualMachineName, newName VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&oldName)
	if err != nil {
		return
	}
	err = c.validateVirtualMachineName(&newName)
	if err != nil {
		return
	}

	// create the change we want to see in the server
	change := brain.VirtualMachine{Name: newName.VirtualMachine}
	if newName.Group != "" || newName.Account != "" {
		// get group
		groupName := GroupName{Group: newName.Group, Account: newName.Account}
		group, groupErr := c.GetGroup(groupName)
		if groupErr != nil {
			return groupErr
		}
		change.GroupID = group.ID
	}

	// PUT the change
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", oldName.Account, oldName.Group, oldName.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.MarshalAndRun(change, nil)
	return err

}

// ReimageVirtualMachine reimages the named virtual machine. This will wipe everything on the first disk in the vm and install a new OS on top of it.
// Note that the machine in question must already be powered off. Once complete, according to the API docs, the vm will be powered on but its autoreboot_on will be false.
func (c *bytemarkClient) ReimageVirtualMachine(name VirtualMachineName, image brain.ImageInstall) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("POST", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s/reimage", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	oldfile := log.LogFile
	log.LogFile = nil
	_, _, err = r.MarshalAndRun(image, nil)
	log.LogFile = oldfile
	return err
}

// ResetVirtualMachine resets the named virtual machine. This is like pressing the reset
// button on a physical computer. This does not cause a new process to be started, so does not apply any pending hardware changes.
// returns nil on success or an error otherwise.
func (c *bytemarkClient) ResetVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("POST", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s/signal", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"signal":"reset"}`), nil)
	return err
}

// RestartVirtualMachine restarts the named virtual machine. This is
// returns nil on success or an error otherwise.
func (c *bytemarkClient) RestartVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"autoreboot_on":true, "power_on": false}`), nil)
	return err
}

// StartVirtualMachine starts the named virtual machine.
// returns nil on success or an error otherwise.
func (c *bytemarkClient) StartVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"autoreboot_on":true, "power_on": true}`), nil)
	return err
}

// StopVirtualMachine starts the named virtual machine.
// returns nil on success or an error otherwise.
func (c *bytemarkClient) StopVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"autoreboot_on":false, "power_on": false}`), nil)
	return err
}

// ShutdownVirtualMachine sends an ACPI shutdown to the VM. This will cause a graceful shutdown of the machine
// returns nil on success or an error otherwise.
func (c *bytemarkClient) ShutdownVirtualMachine(name VirtualMachineName, stayoff bool) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return
	}
	var r *Request
	if stayoff {
		r, err = c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
		if err != nil {
			return
		}

		_, _, err = r.Run(bytes.NewBufferString(`{"autoreboot_on":false}`), nil)
		if err != nil {
			return
		}
	}
	r, err = c.BuildRequest("POST", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s/signal", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"signal": "powerdown"}`), nil)
	return err
}

// UndeleteVirtualMachine changes the deleted flag on a VM back to false.
// Return nil on success, an error otherwise.
func (c *bytemarkClient) UndeleteVirtualMachine(name VirtualMachineName) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.Run(bytes.NewBufferString(`{"deleted":false}`), nil)
	return err
}

// SetVirtualMachineHardwareProfile specifies the hardware profile on a VM. Optionally locks or unlocks h. profile
// Return nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineHardwareProfile(name VirtualMachineName, profile string, locked ...bool) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}
	hwprofileLock := ""
	if len(locked) > 0 {
		hwprofileLock = `, "hardware_profile_locked": false`
		if locked[0] {
			hwprofileLock = `, "hardware_profile_locked": true`
		}
	}
	profileJSON := fmt.Sprintf(`{"hardware_profile": "%s"%s}`, profile, hwprofileLock)

	_, _, err = r.Run(bytes.NewBufferString(profileJSON), nil)
	return err
}

// SetVirtualMachineHardwareProfileLock locks or unlocks the hardware profile of a VM.
// Return nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineHardwareProfileLock(name VirtualMachineName, locked bool) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	lockJSON := `{"hardware_profile_locked": false}`
	if locked {
		lockJSON = `{"hardware_profile_locked": true}`
	}

	_, _, err = r.Run(bytes.NewBufferString(lockJSON), nil)
	return err
}

// SetVirtualMachineMemory sets the RAM available to a virtual machine in megabytes
// Return nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineMemory(name VirtualMachineName, memory int) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	memoryJSON := fmt.Sprintf(`{"memory": %d}`, memory)

	_, _, err = r.Run(bytes.NewBufferString(memoryJSON), nil)
	return err
}

// SetVirtualMachineCores sets the number of CPUs available to a virtual machine
// Return nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineCores(name VirtualMachineName, cores int) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	coresJSON := fmt.Sprintf(`{"cores": %d}`, cores)

	_, _, err = r.Run(bytes.NewBufferString(coresJSON), nil)
	return err
}

// SetVirtualMachineCDROM sets the URL of a CD to attach to a virtual machine. Set url to "" to remove the CD.
// Returns nil on success, an error otherwise.
func (c *bytemarkClient) SetVirtualMachineCDROM(name VirtualMachineName, url string) (err error) {
	err = c.validateVirtualMachineName(&name)
	if err != nil {
		return err
	}
	r, err := c.BuildRequest("PUT", BrainEndpoint, "/accounts/%s/groups/%s/virtual_machines/%s", name.Account, name.Group, name.VirtualMachine)
	if err != nil {
		return err
	}

	_, _, err = r.MarshalAndRun(brain.VirtualMachine{CdromURL: url}, nil)
	return err
}
<|endoftext|>


<commit_before>package shells

import (
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
)

type AbstractShell struct {
}

func (b *AbstractShell) GetFeatures(features *common.FeaturesInfo) {
	features.Artifacts = true
	features.Cache = true
}

func (b *AbstractShell) GetSupportedOptions() []string {
	return []string{"artifacts", "cache", "dependencies"}
}

func (b *AbstractShell) writeCdBuildDir(w ShellWriter, info common.ShellScriptInfo) {
	w.Cd(info.Build.FullProjectDir())
}

func (b *AbstractShell) writeExports(w ShellWriter, info common.ShellScriptInfo) {
	for _, variable := range info.Build.GetAllVariables() {
		w.Variable(variable)
	}
}

func (b *AbstractShell) writeTLSCAInfo(w ShellWriter, build *common.Build, key string) {
	if build.TLSCAChain != "" {
		w.Variable(common.BuildVariable{
			Key:      key,
			Value:    build.TLSCAChain,
			Public:   true,
			Internal: true,
			File:     true,
		})
	}
}

func (b *AbstractShell) writeCloneCmd(w ShellWriter, build *common.Build, projectDir string) {
	w.Notice("Cloning repository...")
	w.RmDir(projectDir)
	w.Command("git", "clone", build.RepoURL, projectDir)
	w.Cd(projectDir)
}

func (b *AbstractShell) writeFetchCmd(w ShellWriter, build *common.Build, projectDir string, gitDir string) {
	w.IfDirectory(gitDir)
	w.Notice("Fetching changes...")
	w.Cd(projectDir)
	w.Command("git", "clean", "-ffdx")
	w.Command("git", "reset", "--hard")
	w.Command("git", "remote", "set-url", "origin", build.RepoURL)
	w.Command("git", "fetch", "origin", "--prune", "+refs/heads/*:refs/remotes/origin/*", "+refs/tags/*:refs/tags/*")
	w.Else()
	b.writeCloneCmd(w, build, projectDir)
	w.EndIf()
}

func (b *AbstractShell) writeCheckoutCmd(w ShellWriter, build *common.Build) {
	w.Notice("Checking out %s as %s...", build.Sha[0:8], build.RefName)
	// We remove a git index file, this is required if `git checkout` is terminated
	w.RmFile(".git/index.lock")
	w.Command("git", "checkout", build.Sha)
}

func (b *AbstractShell) cacheFile(build *common.Build, userKey string) (key, file string) {
	if build.CacheDir == "" {
		return
	}

	// Deduce cache key
	key = path.Join(build.Name, build.RefName)
	if userKey != "" {
		key = build.GetAllVariables().ExpandValue(userKey)
	}

	// Ignore cache without the key
	if key == "" {
		return
	}

	file = path.Join(build.CacheDir, key, "cache.zip")
	file, err := filepath.Rel(build.BuildDir, file)
	if err != nil {
		return "", ""
	}
	return
}

func (o *archivingOptions) CommandArguments() (args []string) {
	for _, path := range o.Paths {
		args = append(args, "--path", path)
	}

	if o.Untracked {
		args = append(args, "--untracked")
	}
	return
}

func (b *AbstractShell) cacheExtractor(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.RunnerCommand == "" {
		w.Warning("The cache is not supported in this executor.")
		return
	}

	// Skip restoring cache if no cache is defined
	if archiverArgs := options.CommandArguments(); len(archiverArgs) == 0 {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-extractor",
		"--file", cacheFile,
	}

	// Generate cache download address
	if url := getCacheDownloadURL(info.Build, cacheKey); url != "" {
		args = append(args, "--url", url)
	}

	// Execute archive command
	w.Notice("Checking cache for %s...", cacheKey)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) downloadArtifacts(w ShellWriter, build *common.BuildInfo, info common.ShellScriptInfo) {
	if info.RunnerCommand == "" {
		w.Warning("The artifacts downloading is not supported in this executor.")
		return
	}

	args := []string{
		"artifacts-downloader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		build.Token,
		"--id",
		strconv.Itoa(build.ID),
	}

	w.Notice("Downloading artifacts for %s (%d)...", build.Name, build.ID)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) downloadAllArtifacts(w ShellWriter, dependencies *dependencies, info common.ShellScriptInfo) {
	for _, otherBuild := range info.Build.DependsOnBuilds {
		if otherBuild.Artifacts == nil || otherBuild.Artifacts.Filename == "" {
			continue
		}
		if !dependencies.IsDependent(otherBuild.Name) {
			continue
		}
		b.downloadArtifacts(w, &otherBuild, info)
	}
}

func (b *AbstractShell) GeneratePreBuild(w ShellWriter, info common.ShellScriptInfo) {
	b.writeExports(w, info)

	build := info.Build
	projectDir := build.FullProjectDir()
	gitDir := path.Join(build.FullProjectDir(), ".git")

	b.writeTLSCAInfo(w, info.Build, "GIT_SSL_CAINFO")
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	if build.AllowGitFetch {
		b.writeFetchCmd(w, build, projectDir, gitDir)
	} else {
		b.writeCloneCmd(w, build, projectDir)
	}

	b.writeCheckoutCmd(w, build)

	// Parse options
	var options shellOptions
	info.Build.Options.Decode(&options)

	// Try to restore from main cache, if not found cache for master
	b.cacheExtractor(w, options.Cache, info)

	// Process all artifacts
	b.downloadAllArtifacts(w, options.Dependencies, info)
}

func (b *AbstractShell) GenerateCommands(w ShellWriter, info common.ShellScriptInfo) {
	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	commands := info.Build.Commands
	commands = strings.TrimSpace(commands)
	for _, command := range strings.Split(commands, "\n") {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}
}

func (b *AbstractShell) cacheArchiver(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.RunnerCommand == "" {
		w.Warning("The cache is not supported in this executor.")
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-archiver",
		"--file", cacheFile,
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Generate cache upload address
	if url := getCacheUploadURL(info.Build, cacheKey); url != "" {
		args = append(args, "--url", url)
	}

	// Execute archive command
	w.Notice("Creating cache %s...", cacheKey)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) uploadArtifacts(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.Build.Runner.URL == "" {
		return
	}
	if info.RunnerCommand == "" {
		w.Warning("The artifacts uploading is not supported in this executor.")
		return
	}

	args := []string{
		"artifacts-uploader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		info.Build.Token,
		"--id",
		strconv.Itoa(info.Build.ID),
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Get artifacts:name
	if name, ok := info.Build.Options.GetString("artifacts", "name"); ok && name != "" {
		args = append(args, "--name", name)
	}

	w.Notice("Uploading artifacts...")
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) GeneratePostBuild(w ShellWriter, info common.ShellScriptInfo) {
	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Parse options
	var options shellOptions
	info.Build.Options.Decode(&options)

	// Find cached files and archive them
	b.cacheArchiver(w, options.Cache, info)

	// Upload artifacts
	b.uploadArtifacts(w, options.Artifacts, info)
}
<commit_msg>Make checkout quiet<commit_after>package shells

import (
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
)

type AbstractShell struct {
}

func (b *AbstractShell) GetFeatures(features *common.FeaturesInfo) {
	features.Artifacts = true
	features.Cache = true
}

func (b *AbstractShell) GetSupportedOptions() []string {
	return []string{"artifacts", "cache", "dependencies"}
}

func (b *AbstractShell) writeCdBuildDir(w ShellWriter, info common.ShellScriptInfo) {
	w.Cd(info.Build.FullProjectDir())
}

func (b *AbstractShell) writeExports(w ShellWriter, info common.ShellScriptInfo) {
	for _, variable := range info.Build.GetAllVariables() {
		w.Variable(variable)
	}
}

func (b *AbstractShell) writeTLSCAInfo(w ShellWriter, build *common.Build, key string) {
	if build.TLSCAChain != "" {
		w.Variable(common.BuildVariable{
			Key:      key,
			Value:    build.TLSCAChain,
			Public:   true,
			Internal: true,
			File:     true,
		})
	}
}

func (b *AbstractShell) writeCloneCmd(w ShellWriter, build *common.Build, projectDir string) {
	w.Notice("Cloning repository...")
	w.RmDir(projectDir)
	w.Command("git", "clone", build.RepoURL, projectDir)
	w.Cd(projectDir)
}

func (b *AbstractShell) writeFetchCmd(w ShellWriter, build *common.Build, projectDir string, gitDir string) {
	w.IfDirectory(gitDir)
	w.Notice("Fetching changes...")
	w.Cd(projectDir)
	w.Command("git", "clean", "-ffdx")
	w.Command("git", "reset", "--hard")
	w.Command("git", "remote", "set-url", "origin", build.RepoURL)
	w.Command("git", "fetch", "origin", "--prune", "+refs/heads/*:refs/remotes/origin/*", "+refs/tags/*:refs/tags/*")
	w.Else()
	b.writeCloneCmd(w, build, projectDir)
	w.EndIf()
}

func (b *AbstractShell) writeCheckoutCmd(w ShellWriter, build *common.Build) {
	w.Notice("Checking out %s as %s...", build.Sha[0:8], build.RefName)
	// We remove a git index file, this is required if `git checkout` is terminated
	w.RmFile(".git/index.lock")
	w.Command("git", "checkout", "-q", build.Sha)
}

func (b *AbstractShell) cacheFile(build *common.Build, userKey string) (key, file string) {
	if build.CacheDir == "" {
		return
	}

	// Deduce cache key
	key = path.Join(build.Name, build.RefName)
	if userKey != "" {
		key = build.GetAllVariables().ExpandValue(userKey)
	}

	// Ignore cache without the key
	if key == "" {
		return
	}

	file = path.Join(build.CacheDir, key, "cache.zip")
	file, err := filepath.Rel(build.BuildDir, file)
	if err != nil {
		return "", ""
	}
	return
}

func (o *archivingOptions) CommandArguments() (args []string) {
	for _, path := range o.Paths {
		args = append(args, "--path", path)
	}

	if o.Untracked {
		args = append(args, "--untracked")
	}
	return
}

func (b *AbstractShell) cacheExtractor(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.RunnerCommand == "" {
		w.Warning("The cache is not supported in this executor.")
		return
	}

	// Skip restoring cache if no cache is defined
	if archiverArgs := options.CommandArguments(); len(archiverArgs) == 0 {
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-extractor",
		"--file", cacheFile,
	}

	// Generate cache download address
	if url := getCacheDownloadURL(info.Build, cacheKey); url != "" {
		args = append(args, "--url", url)
	}

	// Execute archive command
	w.Notice("Checking cache for %s...", cacheKey)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) downloadArtifacts(w ShellWriter, build *common.BuildInfo, info common.ShellScriptInfo) {
	if info.RunnerCommand == "" {
		w.Warning("The artifacts downloading is not supported in this executor.")
		return
	}

	args := []string{
		"artifacts-downloader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		build.Token,
		"--id",
		strconv.Itoa(build.ID),
	}

	w.Notice("Downloading artifacts for %s (%d)...", build.Name, build.ID)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) downloadAllArtifacts(w ShellWriter, dependencies *dependencies, info common.ShellScriptInfo) {
	for _, otherBuild := range info.Build.DependsOnBuilds {
		if otherBuild.Artifacts == nil || otherBuild.Artifacts.Filename == "" {
			continue
		}
		if !dependencies.IsDependent(otherBuild.Name) {
			continue
		}
		b.downloadArtifacts(w, &otherBuild, info)
	}
}

func (b *AbstractShell) GeneratePreBuild(w ShellWriter, info common.ShellScriptInfo) {
	b.writeExports(w, info)

	build := info.Build
	projectDir := build.FullProjectDir()
	gitDir := path.Join(build.FullProjectDir(), ".git")

	b.writeTLSCAInfo(w, info.Build, "GIT_SSL_CAINFO")
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	if build.AllowGitFetch {
		b.writeFetchCmd(w, build, projectDir, gitDir)
	} else {
		b.writeCloneCmd(w, build, projectDir)
	}

	b.writeCheckoutCmd(w, build)

	// Parse options
	var options shellOptions
	info.Build.Options.Decode(&options)

	// Try to restore from main cache, if not found cache for master
	b.cacheExtractor(w, options.Cache, info)

	// Process all artifacts
	b.downloadAllArtifacts(w, options.Dependencies, info)
}

func (b *AbstractShell) GenerateCommands(w ShellWriter, info common.ShellScriptInfo) {
	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)

	commands := info.Build.Commands
	commands = strings.TrimSpace(commands)
	for _, command := range strings.Split(commands, "\n") {
		command = strings.TrimSpace(command)
		if command != "" {
			w.Notice("$ %s", command)
		} else {
			w.EmptyLine()
		}
		w.Line(command)
		w.CheckForErrors()
	}
}

func (b *AbstractShell) cacheArchiver(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.RunnerCommand == "" {
		w.Warning("The cache is not supported in this executor.")
		return
	}

	// Skip archiving if no cache is defined
	cacheKey, cacheFile := b.cacheFile(info.Build, options.Key)
	if cacheKey == "" {
		return
	}

	args := []string{
		"cache-archiver",
		"--file", cacheFile,
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Generate cache upload address
	if url := getCacheUploadURL(info.Build, cacheKey); url != "" {
		args = append(args, "--url", url)
	}

	// Execute archive command
	w.Notice("Creating cache %s...", cacheKey)
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) uploadArtifacts(w ShellWriter, options *archivingOptions, info common.ShellScriptInfo) {
	if options == nil {
		return
	}
	if info.Build.Runner.URL == "" {
		return
	}
	if info.RunnerCommand == "" {
		w.Warning("The artifacts uploading is not supported in this executor.")
		return
	}

	args := []string{
		"artifacts-uploader",
		"--url",
		info.Build.Runner.URL,
		"--token",
		info.Build.Token,
		"--id",
		strconv.Itoa(info.Build.ID),
	}

	// Create list of files to archive
	archiverArgs := options.CommandArguments()
	if len(archiverArgs) == 0 {
		// Skip creating archive
		return
	}
	args = append(args, archiverArgs...)

	// Get artifacts:name
	if name, ok := info.Build.Options.GetString("artifacts", "name"); ok && name != "" {
		args = append(args, "--name", name)
	}

	w.Notice("Uploading artifacts...")
	w.Command(info.RunnerCommand, args...)
}

func (b *AbstractShell) GeneratePostBuild(w ShellWriter, info common.ShellScriptInfo) {
	b.writeExports(w, info)
	b.writeCdBuildDir(w, info)
	b.writeTLSCAInfo(w, info.Build, "CI_SERVER_TLS_CA_FILE")

	// Parse options
	var options shellOptions
	info.Build.Options.Decode(&options)

	// Find cached files and archive them
	b.cacheArchiver(w, options.Cache, info)

	// Upload artifacts
	b.uploadArtifacts(w, options.Artifacts, info)
}
<|endoftext|>


<commit_before>package shodan

import (
	"net/http"
	"strings"
	"testing"

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

func TestClient_GetDNSResolve(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	expectedHostnames := []string{"google.com", "bing.com", "idonotexist.local"}

	mux.HandleFunc(resolvePath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "GET", r.Method)

		hostnames := r.URL.Query().Get("hostnames")
		assert.NotEmpty(t, hostnames)

		splited := strings.Split(hostnames, ",")
		assert.Len(t, splited, len(expectedHostnames))

		w.Write(getStub(t, "dns_resolve"))
	})

	resolve, err := client.GetDNSResolve(expectedHostnames)

	assert.Nil(t, err)
	assert.Len(t, resolve, len(expectedHostnames))

	for _, host := range expectedHostnames {
		_, ok := resolve[host]
		assert.True(t, ok)
	}
}

func TestClient_GetDNSReverse(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	expectedIPs := []string{"74.125.227.244", "92.63.108.40"}

	mux.HandleFunc(reversePath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "GET", r.Method)

		ips := r.URL.Query().Get("ips")
		assert.NotEmpty(t, ips)

		splited := strings.Split(ips, ",")
		assert.Len(t, splited, len(expectedIPs))

		w.Write(getStub(t, "dns_reverse"))
	})

	reversed, err := client.GetDNSReverse(expectedIPs)

	assert.Nil(t, err)
	assert.Len(t, reversed, len(expectedIPs))

	for _, ip := range expectedIPs {
		_, ok := reversed[ip]
		assert.True(t, ok)
	}
}

func TestClient_GetDNSReverse_invalidIP(t *testing.T) {
	_, err := client.GetDNSReverse([]string{"74.125.227", "63.11", "2747393"})

	assert.NotNil(t, err)
}
<commit_msg>check error for type net.ParseError in dns reverse test<commit_after>package shodan

import (
	"net/http"
	"strings"
	"testing"

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

func TestClient_GetDNSResolve(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	expectedHostnames := []string{"google.com", "bing.com", "idonotexist.local"}

	mux.HandleFunc(resolvePath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "GET", r.Method)

		hostnames := r.URL.Query().Get("hostnames")
		assert.NotEmpty(t, hostnames)

		splited := strings.Split(hostnames, ",")
		assert.Len(t, splited, len(expectedHostnames))

		w.Write(getStub(t, "dns_resolve"))
	})

	resolve, err := client.GetDNSResolve(expectedHostnames)

	assert.Nil(t, err)
	assert.Len(t, resolve, len(expectedHostnames))

	for _, host := range expectedHostnames {
		_, ok := resolve[host]
		assert.True(t, ok)
	}
}

func TestClient_GetDNSReverse(t *testing.T) {
	setUpTestServe()
	defer tearDownTestServe()

	expectedIPs := []string{"74.125.227.244", "92.63.108.40"}

	mux.HandleFunc(reversePath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "GET", r.Method)

		ips := r.URL.Query().Get("ips")
		assert.NotEmpty(t, ips)

		splited := strings.Split(ips, ",")
		assert.Len(t, splited, len(expectedIPs))

		w.Write(getStub(t, "dns_reverse"))
	})

	reversed, err := client.GetDNSReverse(expectedIPs)

	assert.Nil(t, err)
	assert.Len(t, reversed, len(expectedIPs))

	for _, ip := range expectedIPs {
		_, ok := reversed[ip]
		assert.True(t, ok)
	}
}

func TestClient_GetDNSReverse_invalidIP(t *testing.T) {
	_, err := client.GetDNSReverse([]string{"74.125.227", "63.11", "2747393"})

	assert.NotNil(t, err)
	_, ok := err.(*net.ParseError)
	assert.True(t, ok)
}
<|endoftext|>