<commit_before>package main

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"
	"unsafe"

	"github.com/Microsoft/go-winio"
	runhcsopts "github.com/Microsoft/hcsshim/cmd/containerd-shim-runhcs-v1/options"
	"github.com/Microsoft/hcsshim/internal/shimdiag"
	"github.com/Microsoft/hcsshim/pkg/octtrpc"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/runtime/v2/task"
	"github.com/containerd/ttrpc"
	"github.com/containerd/typeurl"
	"github.com/gogo/protobuf/proto"
	"github.com/gogo/protobuf/types"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
	"golang.org/x/sys/windows"
)

var svc *service

var serveCommand = cli.Command{
	Name:           "serve",
	Hidden:         true,
	SkipArgReorder: true,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "socket",
			Usage: "the socket path to serve",
		},
		cli.BoolFlag{
			Name:  "is-sandbox",
			Usage: "is the task id a Kubernetes sandbox id",
		},
	},
	Action: func(ctx *cli.Context) error {
		// On Windows the serve command is internally used to actually create
		// the process that hosts the containerd/ttrpc entrypoint to the Runtime
		// V2 API's. The model requires this 2nd invocation of the shim process
		// so that the 1st invocation can return the address via stdout on
		// success.
		//
		// The activation model for this shim is as follows:
		//
		// The public invocation of `shim start` is called which internally
		// decides to either return the address of an existing shim or serve a
		// new one. If serve is decided it execs this entry point `shim serve`.
		// The handoff logic is that this shim will serve the ttrpc entrypoint
		// with only stderr set by the caller. Once the shim has successfully
		// served the entrypoint it is required to close stderr to alert the
		// caller it has completed to the point of handoff. If it fails it will
		// write the error to stderr and the caller will forward the error on as
		// part of the `shim start` failure path. Once successfully served the
		// shim `MUST` not use any std handles. The shim can log any errors to
		// the upstream caller by listening for a log connection and streaming
		// the events.

		var lerrs chan error

		// Default values for shim options.
		shimOpts := &runhcsopts.Options{
			Debug:     false,
			DebugType: runhcsopts.Options_NPIPE,
		}

		// containerd passes the shim options protobuf via stdin.
		newShimOpts, err := readOptions(os.Stdin)
		if err != nil {
			return errors.Wrap(err, "failed to read shim options from stdin")
		} else if newShimOpts != nil {
			// We received a valid shim options struct.
			shimOpts = newShimOpts
		}

		if shimOpts.Debug {
			logrus.SetLevel(logrus.DebugLevel)
		}

		switch shimOpts.DebugType {
		case runhcsopts.Options_NPIPE:
			logrus.SetFormatter(&logrus.TextFormatter{
				TimestampFormat: log.RFC3339NanoFixed,
				FullTimestamp:   true,
			})
			// Setup the log listener
			//
			// TODO: JTERRY75 we need this to be the reconnect log listener or
			// switch to events
			// TODO: JTERRY75 switch containerd to use the protected path.
			//const logAddrFmt = "\\\\.\\pipe\\ProtectedPrefix\\Administrators\\containerd-shim-%s-%s-log"
			const logAddrFmt = "\\\\.\\pipe\\containerd-shim-%s-%s-log"
			logl, err := winio.ListenPipe(fmt.Sprintf(logAddrFmt, namespaceFlag, idFlag), nil)
			if err != nil {
				return err
			}
			defer logl.Close()

			lerrs = make(chan error, 1)
			go func() {
				var cur net.Conn
				for {
					// Listen for log connections in the background
					// We assume that there is always only one client
					// which is containerd. If a new connection is
					// accepted, it means that containerd is restarted.
					// Note that logs generated during containerd restart
					// may be lost.
					new, err := logl.Accept()
					if err != nil {
						lerrs <- err
						return
					}
					if cur != nil {
						cur.Close()
					}
					cur = new
					// Switch the logrus output to here. Note: we wont get this
					// connection until the return from `shim start` so we still
					// havent transitioned the error model yet.
					logrus.SetOutput(cur)
				}
			}()
			// Logrus output will be redirected in the goroutine below that
			// handles the pipe connection.
		case runhcsopts.Options_FILE:
			panic("file log output mode is not supported")
		case runhcsopts.Options_ETW:
			logrus.SetFormatter(nopFormatter{})
			logrus.SetOutput(ioutil.Discard)
		}

		os.Stdin.Close()

		// Force the cli.ErrWriter to be os.Stdout for this. We use stderr for
		// the panic.log attached via start.
		cli.ErrWriter = os.Stdout

		socket := ctx.String("socket")
		if !strings.HasPrefix(socket, `\\.\pipe`) {
			return errors.New("socket is required to be pipe address")
		}

		ttrpcAddress := os.Getenv(ttrpcAddressEnv)
		ttrpcEventPublisher, err := newEventPublisher(ttrpcAddress)

		if err != nil {
			return err
		}

		defer func() {
			if err != nil {
				ttrpcEventPublisher.close()
			}
		}()

		// Setup the ttrpc server
		svc = &service{
			events:    ttrpcEventPublisher,
			tid:       idFlag,
			isSandbox: ctx.Bool("is-sandbox"),
		}
		s, err := ttrpc.NewServer(ttrpc.WithUnaryServerInterceptor(octtrpc.ServerInterceptor()))
		if err != nil {
			return err
		}
		defer s.Close()
		task.RegisterTaskService(s, svc)
		shimdiag.RegisterShimDiagService(s, svc)

		sl, err := winio.ListenPipe(socket, nil)
		if err != nil {
			return err
		}
		defer sl.Close()

		serrs := make(chan error, 1)
		defer close(serrs)
		go func() {
			// TODO: JTERRY75 We should use a real context with cancellation shared by
			// the service for shim shutdown gracefully.
			ctx := context.Background()
			if err := trapClosedConnErr(s.Serve(ctx, sl)); err != nil {
				logrus.WithError(err).Fatal("containerd-shim: ttrpc server failure")
				serrs <- err
				return
			}
			serrs <- nil
		}()

		select {
		case err := <-lerrs:
			return err
		case err := <-serrs:
			return err
		case <-time.After(2 * time.Millisecond):
			// TODO: JTERRY75 this is terrible code. Contribue a change to
			// ttrpc that you can:
			//
			// go func () { errs <- s.Serve() }
			// select {
			// case <-errs:
			// case <-s.Ready():
			// }

			// This is our best indication that we have not errored on creation
			// and are successfully serving the API.
			os.Stdout.Close()
		}

		// Wait for the serve API to be shut down.
		<-serrs
		return nil
	},
}

func trapClosedConnErr(err error) error {
	if err == nil || strings.Contains(err.Error(), "use of closed network connection") {
		return nil
	}
	return err
}

// readOptions reads in bytes from the reader and converts it to a shim options
// struct. If no data is available from the reader, returns (nil, nil).
func readOptions(r io.Reader) (*runhcsopts.Options, error) {
	d, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, errors.Wrap(err, "failed to read input")
	}
	if len(d) > 0 {
		var a types.Any
		if err := proto.Unmarshal(d, &a); err != nil {
			return nil, errors.Wrap(err, "failed unmarshaling into Any")
		}
		v, err := typeurl.UnmarshalAny(&a)
		if err != nil {
			return nil, errors.Wrap(err, "failed unmarshaling by typeurl")
		}
		return v.(*runhcsopts.Options), nil
	}
	return nil, nil
}

// createEvent creates a Windows event ACL'd to builtin administrator
// and local system. Can use docker-signal to signal the event.
func createEvent(event string) (windows.Handle, error) {
	ev, _ := windows.UTF16PtrFromString(event)
	sd, err := windows.SecurityDescriptorFromString("D:P(A;;GA;;;BA)(A;;GA;;;SY)")
	if err != nil {
		return 0, errors.Wrapf(err, "failed to get security descriptor for event '%s'", event)
	}
	var sa windows.SecurityAttributes
	sa.Length = uint32(unsafe.Sizeof(sa))
	sa.InheritHandle = 1
	sa.SecurityDescriptor = sd
	h, err := windows.CreateEvent(&sa, 0, 0, ev)
	if h == 0 || err != nil {
		return 0, errors.Wrapf(err, "failed to create event '%s'", event)
	}
	return h, nil
}

// setupDebuggerEvent listens for an event to allow a debugger such as delve
// to attach for advanced debugging. It's called when handling a ContainerCreate
func setupDebuggerEvent() {
	if os.Getenv("CONTAINERD_SHIM_RUNHCS_V1_WAIT_DEBUGGER") == "" {
		return
	}
	event := "Global\\debugger-" + fmt.Sprint(os.Getpid())
	handle, err := createEvent(event)
	if err != nil {
		return
	}
	logrus.WithField("event", event).Info("Halting until signalled")
	_, _ = windows.WaitForSingleObject(handle, windows.INFINITE)
}
<commit_msg>Fix stderr comment in containerd-shim serve command<commit_after>package main

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"
	"unsafe"

	"github.com/Microsoft/go-winio"
	runhcsopts "github.com/Microsoft/hcsshim/cmd/containerd-shim-runhcs-v1/options"
	"github.com/Microsoft/hcsshim/internal/shimdiag"
	"github.com/Microsoft/hcsshim/pkg/octtrpc"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/runtime/v2/task"
	"github.com/containerd/ttrpc"
	"github.com/containerd/typeurl"
	"github.com/gogo/protobuf/proto"
	"github.com/gogo/protobuf/types"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
	"golang.org/x/sys/windows"
)

var svc *service

var serveCommand = cli.Command{
	Name:           "serve",
	Hidden:         true,
	SkipArgReorder: true,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "socket",
			Usage: "the socket path to serve",
		},
		cli.BoolFlag{
			Name:  "is-sandbox",
			Usage: "is the task id a Kubernetes sandbox id",
		},
	},
	Action: func(ctx *cli.Context) error {
		// On Windows the serve command is internally used to actually create
		// the process that hosts the containerd/ttrpc entrypoint to the Runtime
		// V2 API's. The model requires this 2nd invocation of the shim process
		// so that the 1st invocation can return the address via stdout on
		// success.
		//
		// The activation model for this shim is as follows:
		//
		// The public invocation of `shim start` is called which internally
		// decides to either return the address of an existing shim or serve a
		// new one. If serve is decided it execs this entry point `shim serve`.
		// The handoff logic is that this shim will serve the ttrpc entrypoint
		// with only stderr set by the caller. Once the shim has successfully
		// served the entrypoint it is required to close stdout to alert the
		// caller it has completed to the point of handoff. If it fails it will
		// write the error to stderr and the caller will forward the error on as
		// part of the `shim start` failure path. Once successfully served the
		// shim `MUST` not use any std handles. The shim can log any errors to
		// the upstream caller by listening for a log connection and streaming
		// the events.

		var lerrs chan error

		// Default values for shim options.
		shimOpts := &runhcsopts.Options{
			Debug:     false,
			DebugType: runhcsopts.Options_NPIPE,
		}

		// containerd passes the shim options protobuf via stdin.
		newShimOpts, err := readOptions(os.Stdin)
		if err != nil {
			return errors.Wrap(err, "failed to read shim options from stdin")
		} else if newShimOpts != nil {
			// We received a valid shim options struct.
			shimOpts = newShimOpts
		}

		if shimOpts.Debug {
			logrus.SetLevel(logrus.DebugLevel)
		}

		switch shimOpts.DebugType {
		case runhcsopts.Options_NPIPE:
			logrus.SetFormatter(&logrus.TextFormatter{
				TimestampFormat: log.RFC3339NanoFixed,
				FullTimestamp:   true,
			})
			// Setup the log listener
			//
			// TODO: JTERRY75 we need this to be the reconnect log listener or
			// switch to events
			// TODO: JTERRY75 switch containerd to use the protected path.
			//const logAddrFmt = "\\\\.\\pipe\\ProtectedPrefix\\Administrators\\containerd-shim-%s-%s-log"
			const logAddrFmt = "\\\\.\\pipe\\containerd-shim-%s-%s-log"
			logl, err := winio.ListenPipe(fmt.Sprintf(logAddrFmt, namespaceFlag, idFlag), nil)
			if err != nil {
				return err
			}
			defer logl.Close()

			lerrs = make(chan error, 1)
			go func() {
				var cur net.Conn
				for {
					// Listen for log connections in the background
					// We assume that there is always only one client
					// which is containerd. If a new connection is
					// accepted, it means that containerd is restarted.
					// Note that logs generated during containerd restart
					// may be lost.
					new, err := logl.Accept()
					if err != nil {
						lerrs <- err
						return
					}
					if cur != nil {
						cur.Close()
					}
					cur = new
					// Switch the logrus output to here. Note: we wont get this
					// connection until the return from `shim start` so we still
					// havent transitioned the error model yet.
					logrus.SetOutput(cur)
				}
			}()
			// Logrus output will be redirected in the goroutine below that
			// handles the pipe connection.
		case runhcsopts.Options_FILE:
			panic("file log output mode is not supported")
		case runhcsopts.Options_ETW:
			logrus.SetFormatter(nopFormatter{})
			logrus.SetOutput(ioutil.Discard)
		}

		os.Stdin.Close()

		// Force the cli.ErrWriter to be os.Stdout for this. We use stderr for
		// the panic.log attached via start.
		cli.ErrWriter = os.Stdout

		socket := ctx.String("socket")
		if !strings.HasPrefix(socket, `\\.\pipe`) {
			return errors.New("socket is required to be pipe address")
		}

		ttrpcAddress := os.Getenv(ttrpcAddressEnv)
		ttrpcEventPublisher, err := newEventPublisher(ttrpcAddress)

		if err != nil {
			return err
		}

		defer func() {
			if err != nil {
				ttrpcEventPublisher.close()
			}
		}()

		// Setup the ttrpc server
		svc = &service{
			events:    ttrpcEventPublisher,
			tid:       idFlag,
			isSandbox: ctx.Bool("is-sandbox"),
		}
		s, err := ttrpc.NewServer(ttrpc.WithUnaryServerInterceptor(octtrpc.ServerInterceptor()))
		if err != nil {
			return err
		}
		defer s.Close()
		task.RegisterTaskService(s, svc)
		shimdiag.RegisterShimDiagService(s, svc)

		sl, err := winio.ListenPipe(socket, nil)
		if err != nil {
			return err
		}
		defer sl.Close()

		serrs := make(chan error, 1)
		defer close(serrs)
		go func() {
			// TODO: JTERRY75 We should use a real context with cancellation shared by
			// the service for shim shutdown gracefully.
			ctx := context.Background()
			if err := trapClosedConnErr(s.Serve(ctx, sl)); err != nil {
				logrus.WithError(err).Fatal("containerd-shim: ttrpc server failure")
				serrs <- err
				return
			}
			serrs <- nil
		}()

		select {
		case err := <-lerrs:
			return err
		case err := <-serrs:
			return err
		case <-time.After(2 * time.Millisecond):
			// TODO: JTERRY75 this is terrible code. Contribue a change to
			// ttrpc that you can:
			//
			// go func () { errs <- s.Serve() }
			// select {
			// case <-errs:
			// case <-s.Ready():
			// }

			// This is our best indication that we have not errored on creation
			// and are successfully serving the API.
			os.Stdout.Close()
		}

		// Wait for the serve API to be shut down.
		<-serrs
		return nil
	},
}

func trapClosedConnErr(err error) error {
	if err == nil || strings.Contains(err.Error(), "use of closed network connection") {
		return nil
	}
	return err
}

// readOptions reads in bytes from the reader and converts it to a shim options
// struct. If no data is available from the reader, returns (nil, nil).
func readOptions(r io.Reader) (*runhcsopts.Options, error) {
	d, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, errors.Wrap(err, "failed to read input")
	}
	if len(d) > 0 {
		var a types.Any
		if err := proto.Unmarshal(d, &a); err != nil {
			return nil, errors.Wrap(err, "failed unmarshaling into Any")
		}
		v, err := typeurl.UnmarshalAny(&a)
		if err != nil {
			return nil, errors.Wrap(err, "failed unmarshaling by typeurl")
		}
		return v.(*runhcsopts.Options), nil
	}
	return nil, nil
}

// createEvent creates a Windows event ACL'd to builtin administrator
// and local system. Can use docker-signal to signal the event.
func createEvent(event string) (windows.Handle, error) {
	ev, _ := windows.UTF16PtrFromString(event)
	sd, err := windows.SecurityDescriptorFromString("D:P(A;;GA;;;BA)(A;;GA;;;SY)")
	if err != nil {
		return 0, errors.Wrapf(err, "failed to get security descriptor for event '%s'", event)
	}
	var sa windows.SecurityAttributes
	sa.Length = uint32(unsafe.Sizeof(sa))
	sa.InheritHandle = 1
	sa.SecurityDescriptor = sd
	h, err := windows.CreateEvent(&sa, 0, 0, ev)
	if h == 0 || err != nil {
		return 0, errors.Wrapf(err, "failed to create event '%s'", event)
	}
	return h, nil
}

// setupDebuggerEvent listens for an event to allow a debugger such as delve
// to attach for advanced debugging. It's called when handling a ContainerCreate
func setupDebuggerEvent() {
	if os.Getenv("CONTAINERD_SHIM_RUNHCS_V1_WAIT_DEBUGGER") == "" {
		return
	}
	event := "Global\\debugger-" + fmt.Sprint(os.Getpid())
	handle, err := createEvent(event)
	if err != nil {
		return
	}
	logrus.WithField("event", event).Info("Halting until signalled")
	_, _ = windows.WaitForSingleObject(handle, windows.INFINITE)
}
<|endoftext|>


<commit_before>package imgio

import (
	"bytes"
	"crypto/md5"
	"crypto/rand"
	"image"
	"io"
	"testing"

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

func Test_ImageGroup_ReadWriteHash_OneImage(t *testing.T) {
	group := &ImageGroup{
		images: []*rwImage{
			{
				img: image.NewRGBA(image.Rect(0, 0, 100, 100)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 100),
					cursor: 0,
				},
				prw: GentlePoint16ReadWriter{},
			},
		},
	}
	hasher := md5.New()
	buff := bytes.NewBuffer(nil)
	n, err := buff.ReadFrom(io.TeeReader(io.LimitReader(rand.Reader, group.Size()), hasher))
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	firstSum := hasher.Sum(nil)
	hasher.Reset()
	n, err = buff.WriteTo(group)
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	group.Rewind()
	n, err = io.Copy(hasher, group)
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	secondSum := hasher.Sum(nil)
	require.Equal(t, firstSum, secondSum)
}

func Test_ImageGroup_ReadWriteHash_ManyImage(t *testing.T) {
	group := &ImageGroup{
		images: []*rwImage{
			{
				img: image.NewRGBA(image.Rect(0, 0, 100, 100)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 100),
					cursor: 0,
				},
				prw: GentlePoint16ReadWriter{},
			},
			{
				img: image.NewRGBA(image.Rect(0, 0, 100, 100)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 100),
					cursor: 0,
				},
				prw: SimplePoint32ReadWriter{},
			},
			{
				img: image.NewRGBA64(image.Rect(0, 0, 100, 100)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 100),
					cursor: 0,
				},
				prw: SimplePoint64ReadWriter{},
			},
			{
				img: image.NewRGBA(image.Rect(0, 0, 100, 10)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 10),
					cursor: 0,
				},
				prw: SimplePoint32ReadWriter{},
			},
		},
	}
	hasher := md5.New()
	buff := bytes.NewBuffer(nil)
	n, err := buff.ReadFrom(io.TeeReader(io.LimitReader(rand.Reader, group.Size()), hasher))
	require.Nil(t, err)
	require.Equal(t, group.Size(), int64(n))
	firstSum := hasher.Sum(nil)
	hasher.Reset()
	n, err = buff.WriteTo(group)
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	group.Rewind()
	n, err = io.Copy(hasher, group)
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	secondSum := hasher.Sum(nil)
	require.Equal(t, firstSum, secondSum)
}
<commit_msg>Create test case for image group test<commit_after>package imgio

import (
	"bytes"
	"crypto/md5"
	"crypto/rand"
	"image"
	"io"
	"testing"

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

func Test_ImageGroup_ReadWriteHash_OneImage(t *testing.T) {
	group := &ImageGroup{
		images: []*rwImage{
			{
				img: image.NewRGBA(image.Rect(0, 0, 100, 100)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 100),
					cursor: 0,
				},
				prw: GentlePoint16ReadWriter{},
			},
		},
	}
	hasher := md5.New()
	buff := bytes.NewBuffer(nil)
	n, err := buff.ReadFrom(io.TeeReader(io.LimitReader(rand.Reader, group.Size()), hasher))
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	firstSum := hasher.Sum(nil)
	hasher.Reset()
	n, err = buff.WriteTo(group)
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	group.Rewind()
	n, err = io.Copy(hasher, group)
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	secondSum := hasher.Sum(nil)
	require.Equal(t, firstSum, secondSum)
}

func Test_ImageGroup_ReadWriteHash_ManyImage(t *testing.T) {
	group := &ImageGroup{
		images: []*rwImage{
			{
				img: image.NewRGBA(image.Rect(0, 0, 100, 100)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 100),
					cursor: 0,
				},
				prw: GentlePoint16ReadWriter{},
			},
			{
				img: image.NewRGBA(image.Rect(0, 0, 100, 100)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 100),
					cursor: 0,
				},
				prw: SimplePoint32ReadWriter{},
			},
			{
				img: image.NewRGBA64(image.Rect(0, 0, 100, 100)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 100),
					cursor: 0,
				},
				prw: SimplePoint64ReadWriter{},
			},
			{
				img: image.NewRGBA(image.Rect(0, 0, 100, 10)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 10),
					cursor: 0,
				},
				prw: SimplePoint32ReadWriter{},
			},
			{
				img: image.NewRGBA64(image.Rect(0, 0, 100, 52)),
				gen: &SimplePointsSequenceGenerator{
					rect:   image.Rect(0, 0, 100, 52),
					cursor: 0,
				},
				prw: SimplePoint64ReadWriter{},
			},
		},
	}
	hasher := md5.New()
	buff := bytes.NewBuffer(nil)
	n, err := buff.ReadFrom(io.TeeReader(io.LimitReader(rand.Reader, group.Size()), hasher))
	require.Nil(t, err)
	require.Equal(t, group.Size(), int64(n))
	firstSum := hasher.Sum(nil)
	hasher.Reset()
	n, err = buff.WriteTo(group)
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	group.Rewind()
	n, err = io.Copy(hasher, group)
	require.Nil(t, err)
	require.Equal(t, group.Size(), n)
	secondSum := hasher.Sum(nil)
	require.Equal(t, firstSum, secondSum)
}
<|endoftext|>


<commit_before>// Package knn implements a K Nearest Neighbors object, capable of both classification
// and regression. It accepts data in the form of a slice of float64s, which are then reshaped
// into a X by Y matrix.
package knn

import (
	"errors"
	"fmt"

	"github.com/gonum/matrix"
	"github.com/gonum/matrix/mat64"
	"github.com/sjwhitworth/golearn/base"
	"github.com/sjwhitworth/golearn/kdtree"
	"github.com/sjwhitworth/golearn/metrics/pairwise"
	"github.com/sjwhitworth/golearn/utilities"
)

// A KNNClassifier consists of a data matrix, associated labels in the same order as the matrix, searching algorithm, and a distance function.
// The accepted distance functions at this time are 'euclidean', 'manhattan', and 'cosine'.
// The accepted searching algorithm here are 'linear', and 'kdtree'.
// Optimisations only occur when things are identically group into identical
// AttributeGroups, which don't include the class variable, in the same order.
// Using weighted KNN when Weighted set to be true (default: false).
type KNNClassifier struct {
	base.BaseEstimator
	TrainingData       base.FixedDataGrid
	DistanceFunc       string
	Algorithm          string
	NearestNeighbours  int
	AllowOptimisations bool
	Weighted           bool
}

// NewKnnClassifier returns a new classifier
func NewKnnClassifier(distfunc, algorithm string, neighbours int) *KNNClassifier {
	KNN := KNNClassifier{}
	KNN.DistanceFunc = distfunc
	KNN.Algorithm = algorithm
	KNN.NearestNeighbours = neighbours
	KNN.Weighted = false
	KNN.AllowOptimisations = true
	return &KNN
}

// Fit stores the training data for later
func (KNN *KNNClassifier) Fit(trainingData base.FixedDataGrid) error {
	KNN.TrainingData = trainingData
	return nil
}

func (KNN *KNNClassifier) canUseOptimisations(what base.FixedDataGrid) bool {
	// Check that the two have exactly the same layout
	if !base.CheckStrictlyCompatible(what, KNN.TrainingData) {
		return false
	}
	// Check that the two are DenseInstances
	whatd, ok1 := what.(*base.DenseInstances)
	_, ok2 := KNN.TrainingData.(*base.DenseInstances)
	if !ok1 || !ok2 {
		return false
	}
	// Check that no Class Attributes are mixed in with the data
	classAttrs := whatd.AllClassAttributes()
	normalAttrs := base.NonClassAttributes(whatd)
	// Retrieve all the AGs
	ags := whatd.AllAttributeGroups()
	classAttrGroups := make([]base.AttributeGroup, 0)
	for agName := range ags {
		ag := ags[agName]
		attrs := ag.Attributes()
		matched := false
		for _, a := range attrs {
			for _, c := range classAttrs {
				if a.Equals(c) {
					matched = true
				}
			}
		}
		if matched {
			classAttrGroups = append(classAttrGroups, ag)
		}
	}
	for _, cag := range classAttrGroups {
		attrs := cag.Attributes()
		common := base.AttributeIntersect(normalAttrs, attrs)
		if len(common) != 0 {
			return false
		}
	}

	// Check that all of the Attributes are numeric
	for _, a := range normalAttrs {
		if _, ok := a.(*base.FloatAttribute); !ok {
			return false
		}
	}
	// If that's fine, return true
	return true
}

// Predict returns a classification for the vector, based on a vector input, using the KNN algorithm.
func (KNN *KNNClassifier) Predict(what base.FixedDataGrid) (base.FixedDataGrid, error) {
	// Check what distance function we are using
	var distanceFunc pairwise.PairwiseDistanceFunc
	switch KNN.DistanceFunc {
	case "euclidean":
		distanceFunc = pairwise.NewEuclidean()
	case "manhattan":
		distanceFunc = pairwise.NewManhattan()
	case "cosine":
		distanceFunc = pairwise.NewCosine()
	default:
		return nil, errors.New("unsupported distance function")
	}

	// Check what searching algorith, we are using
	if KNN.Algorithm != "linear" && KNN.Algorithm != "kdtree" {
		return nil, errors.New("unsupported searching algorithm")
	}

	// Check Compatibility
	allAttrs := base.CheckCompatible(what, KNN.TrainingData)
	if allAttrs == nil {
		// Don't have the same Attributes
		return nil, errors.New("attributes not compatible")
	}

	// Use optimised version if permitted
	if KNN.Algorithm == "linear" && KNN.AllowOptimisations {
		if KNN.DistanceFunc == "euclidean" {
			if KNN.canUseOptimisations(what) {
				return KNN.optimisedEuclideanPredict(what.(*base.DenseInstances)), nil
			}
		}
	}
	fmt.Println("Optimisations are switched off")

	// Remove the Attributes which aren't numeric
	allNumericAttrs := make([]base.Attribute, 0)
	for _, a := range allAttrs {
		if fAttr, ok := a.(*base.FloatAttribute); ok {
			allNumericAttrs = append(allNumericAttrs, fAttr)
		}
	}

	// If every Attribute is a FloatAttribute, then we remove the last one
	// because that is the Attribute we are trying to predict.
	if len(allNumericAttrs) == len(allAttrs) {
		allNumericAttrs = allNumericAttrs[:len(allNumericAttrs)-1]
	}

	// Generate return vector
	ret := base.GeneratePredictionVector(what)

	// Resolve Attribute specifications for both
	whatAttrSpecs := base.ResolveAttributes(what, allNumericAttrs)
	trainAttrSpecs := base.ResolveAttributes(KNN.TrainingData, allNumericAttrs)

	// Reserve storage for most the most similar items
	distances := make(map[int]float64)

	// Reserve storage for voting map
	maxmapInt := make(map[string]int)
	maxmapFloat := make(map[string]float64)

	// Reserve storage for row computations
	trainRowBuf := make([]float64, len(allNumericAttrs))
	predRowBuf := make([]float64, len(allNumericAttrs))

	_, maxRow := what.Size()
	curRow := 0

	// build kdtree if algorithm is 'kdtree'
	kd := kdtree.New()
	srcRowNoMap := make([]int, 0)
	if KNN.Algorithm == "kdtree" {
		buildData := make([][]float64, 0)
		KNN.TrainingData.MapOverRows(trainAttrSpecs, func(trainRow [][]byte, srcRowNo int) (bool, error) {
			oneData := make([]float64, len(allNumericAttrs))
			// Read the float values out
			for i, _ := range allNumericAttrs {
				oneData[i] = base.UnpackBytesToFloat(trainRow[i])
			}
			srcRowNoMap = append(srcRowNoMap, srcRowNo)
			buildData = append(buildData, oneData)
			return true, nil
		})

		err := kd.Build(buildData)
		if err != nil {
			return nil, err
		}
	}

	// Iterate over all outer rows
	what.MapOverRows(whatAttrSpecs, func(predRow [][]byte, predRowNo int) (bool, error) {

		if (curRow%1) == 0 && curRow > 0 {
			fmt.Printf("KNN: %.2f %% done\n", float64(curRow)*100.0/float64(maxRow))
		}
		curRow++

		// Read the float values out
		for i, _ := range allNumericAttrs {
			predRowBuf[i] = base.UnpackBytesToFloat(predRow[i])
		}

		predMat := utilities.FloatsToMatrix(predRowBuf)

		switch KNN.Algorithm {
		case "linear":
			// Find the closest match in the training data
			KNN.TrainingData.MapOverRows(trainAttrSpecs, func(trainRow [][]byte, srcRowNo int) (bool, error) {
				// Read the float values out
				for i, _ := range allNumericAttrs {
					trainRowBuf[i] = base.UnpackBytesToFloat(trainRow[i])
				}

				// Compute the distance
				trainMat := utilities.FloatsToMatrix(trainRowBuf)
				distances[srcRowNo] = distanceFunc.Distance(predMat, trainMat)
				return true, nil
			})

			sorted := utilities.SortIntMap(distances)
			values := sorted[:KNN.NearestNeighbours]

			length := make([]float64, KNN.NearestNeighbours)
			for k, v := range values {
				length[k] = distances[v]
			}

			var maxClass string
			if KNN.Weighted {
				maxClass = KNN.weightedVote(maxmapFloat, values, length)
			} else {
				maxClass = KNN.vote(maxmapInt, values)
			}
			base.SetClass(ret, predRowNo, maxClass)

		case "kdtree":
			// search kdtree
			values, length, err := kd.Search(KNN.NearestNeighbours, distanceFunc, predRowBuf)
			if err != nil {
				return false, err
			}

			// map values to srcRowNo
			for k, v := range values {
				values[k] = srcRowNoMap[v]
			}

			var maxClass string
			if KNN.Weighted {
				maxClass = KNN.weightedVote(maxmapFloat, values, length)
			} else {
				maxClass = KNN.vote(maxmapInt, values)
			}
			base.SetClass(ret, predRowNo, maxClass)
		}

		return true, nil

	})

	return ret, nil
}

func (KNN *KNNClassifier) String() string {
	return fmt.Sprintf("KNNClassifier(%s, %d)", KNN.DistanceFunc, KNN.NearestNeighbours)
}

func (KNN *KNNClassifier) vote(maxmap map[string]int, values []int) string {
	// Reset maxMap
	for a := range maxmap {
		maxmap[a] = 0
	}

	// Refresh maxMap
	for _, elem := range values {
		label := base.GetClass(KNN.TrainingData, elem)
		if _, ok := maxmap[label]; ok {
			maxmap[label]++
		} else {
			maxmap[label] = 1
		}
	}

	// Sort the maxMap
	var maxClass string
	maxVal := -1
	for a := range maxmap {
		if maxmap[a] > maxVal {
			maxVal = maxmap[a]
			maxClass = a
		}
	}
	return maxClass
}

func (KNN *KNNClassifier) weightedVote(maxmap map[string]float64, values []int, length []float64) string {
	// Reset maxMap
	for a := range maxmap {
		maxmap[a] = 0
	}

	// Refresh maxMap
	for k, elem := range values {
		label := base.GetClass(KNN.TrainingData, elem)
		if _, ok := maxmap[label]; ok {
			maxmap[label] += (1 / length[k])
		} else {
			maxmap[label] = (1 / length[k])
		}
	}

	// Sort the maxMap
	var maxClass string
	maxVal := -1.0
	for a := range maxmap {
		if maxmap[a] > maxVal {
			maxVal = maxmap[a]
			maxClass = a
		}
	}
	return maxClass
}

// A KNNRegressor consists of a data matrix, associated result variables in the same order as the matrix, and a name.
type KNNRegressor struct {
	base.BaseEstimator
	Values       []float64
	DistanceFunc string
}

// NewKnnRegressor mints a new classifier.
func NewKnnRegressor(distfunc string) *KNNRegressor {
	KNN := KNNRegressor{}
	KNN.DistanceFunc = distfunc
	return &KNN
}

func (KNN *KNNRegressor) Fit(values []float64, numbers []float64, rows int, cols int) {
	if rows != len(values) {
		panic(matrix.ErrShape)
	}

	KNN.Data = mat64.NewDense(rows, cols, numbers)
	KNN.Values = values
}

func (KNN *KNNRegressor) Predict(vector *mat64.Dense, K int) float64 {
	// Get the number of rows
	rows, _ := KNN.Data.Dims()
	rownumbers := make(map[int]float64)
	labels := make([]float64, 0)

	// Check what distance function we are using
	var distanceFunc pairwise.PairwiseDistanceFunc
	switch KNN.DistanceFunc {
	case "euclidean":
		distanceFunc = pairwise.NewEuclidean()
	case "manhattan":
		distanceFunc = pairwise.NewManhattan()
	default:
		panic("unsupported distance function")
	}

	for i := 0; i < rows; i++ {
		row := KNN.Data.RowView(i)
		distance := distanceFunc.Distance(utilities.VectorToMatrix(row), vector)
		rownumbers[i] = distance
	}

	sorted := utilities.SortIntMap(rownumbers)
	values := sorted[:K]

	var sum float64
	for _, elem := range values {
		value := KNN.Values[elem]
		labels = append(labels, value)
		sum += value
	}

	average := sum / float64(K)
	return average
}
<commit_msg>replace output \n with \r when training knn<commit_after>// Package knn implements a K Nearest Neighbors object, capable of both classification
// and regression. It accepts data in the form of a slice of float64s, which are then reshaped
// into a X by Y matrix.
package knn

import (
	"errors"
	"fmt"

	"github.com/gonum/matrix"
	"github.com/gonum/matrix/mat64"
	"github.com/sjwhitworth/golearn/base"
	"github.com/sjwhitworth/golearn/kdtree"
	"github.com/sjwhitworth/golearn/metrics/pairwise"
	"github.com/sjwhitworth/golearn/utilities"
)

// A KNNClassifier consists of a data matrix, associated labels in the same order as the matrix, searching algorithm, and a distance function.
// The accepted distance functions at this time are 'euclidean', 'manhattan', and 'cosine'.
// The accepted searching algorithm here are 'linear', and 'kdtree'.
// Optimisations only occur when things are identically group into identical
// AttributeGroups, which don't include the class variable, in the same order.
// Using weighted KNN when Weighted set to be true (default: false).
type KNNClassifier struct {
	base.BaseEstimator
	TrainingData       base.FixedDataGrid
	DistanceFunc       string
	Algorithm          string
	NearestNeighbours  int
	AllowOptimisations bool
	Weighted           bool
}

// NewKnnClassifier returns a new classifier
func NewKnnClassifier(distfunc, algorithm string, neighbours int) *KNNClassifier {
	KNN := KNNClassifier{}
	KNN.DistanceFunc = distfunc
	KNN.Algorithm = algorithm
	KNN.NearestNeighbours = neighbours
	KNN.Weighted = false
	KNN.AllowOptimisations = true
	return &KNN
}

// Fit stores the training data for later
func (KNN *KNNClassifier) Fit(trainingData base.FixedDataGrid) error {
	KNN.TrainingData = trainingData
	return nil
}

func (KNN *KNNClassifier) canUseOptimisations(what base.FixedDataGrid) bool {
	// Check that the two have exactly the same layout
	if !base.CheckStrictlyCompatible(what, KNN.TrainingData) {
		return false
	}
	// Check that the two are DenseInstances
	whatd, ok1 := what.(*base.DenseInstances)
	_, ok2 := KNN.TrainingData.(*base.DenseInstances)
	if !ok1 || !ok2 {
		return false
	}
	// Check that no Class Attributes are mixed in with the data
	classAttrs := whatd.AllClassAttributes()
	normalAttrs := base.NonClassAttributes(whatd)
	// Retrieve all the AGs
	ags := whatd.AllAttributeGroups()
	classAttrGroups := make([]base.AttributeGroup, 0)
	for agName := range ags {
		ag := ags[agName]
		attrs := ag.Attributes()
		matched := false
		for _, a := range attrs {
			for _, c := range classAttrs {
				if a.Equals(c) {
					matched = true
				}
			}
		}
		if matched {
			classAttrGroups = append(classAttrGroups, ag)
		}
	}
	for _, cag := range classAttrGroups {
		attrs := cag.Attributes()
		common := base.AttributeIntersect(normalAttrs, attrs)
		if len(common) != 0 {
			return false
		}
	}

	// Check that all of the Attributes are numeric
	for _, a := range normalAttrs {
		if _, ok := a.(*base.FloatAttribute); !ok {
			return false
		}
	}
	// If that's fine, return true
	return true
}

// Predict returns a classification for the vector, based on a vector input, using the KNN algorithm.
func (KNN *KNNClassifier) Predict(what base.FixedDataGrid) (base.FixedDataGrid, error) {
	// Check what distance function we are using
	var distanceFunc pairwise.PairwiseDistanceFunc
	switch KNN.DistanceFunc {
	case "euclidean":
		distanceFunc = pairwise.NewEuclidean()
	case "manhattan":
		distanceFunc = pairwise.NewManhattan()
	case "cosine":
		distanceFunc = pairwise.NewCosine()
	default:
		return nil, errors.New("unsupported distance function")
	}

	// Check what searching algorith, we are using
	if KNN.Algorithm != "linear" && KNN.Algorithm != "kdtree" {
		return nil, errors.New("unsupported searching algorithm")
	}

	// Check Compatibility
	allAttrs := base.CheckCompatible(what, KNN.TrainingData)
	if allAttrs == nil {
		// Don't have the same Attributes
		return nil, errors.New("attributes not compatible")
	}

	// Use optimised version if permitted
	if KNN.Algorithm == "linear" && KNN.AllowOptimisations {
		if KNN.DistanceFunc == "euclidean" {
			if KNN.canUseOptimisations(what) {
				return KNN.optimisedEuclideanPredict(what.(*base.DenseInstances)), nil
			}
		}
	}
	fmt.Println("Optimisations are switched off")

	// Remove the Attributes which aren't numeric
	allNumericAttrs := make([]base.Attribute, 0)
	for _, a := range allAttrs {
		if fAttr, ok := a.(*base.FloatAttribute); ok {
			allNumericAttrs = append(allNumericAttrs, fAttr)
		}
	}

	// If every Attribute is a FloatAttribute, then we remove the last one
	// because that is the Attribute we are trying to predict.
	if len(allNumericAttrs) == len(allAttrs) {
		allNumericAttrs = allNumericAttrs[:len(allNumericAttrs)-1]
	}

	// Generate return vector
	ret := base.GeneratePredictionVector(what)

	// Resolve Attribute specifications for both
	whatAttrSpecs := base.ResolveAttributes(what, allNumericAttrs)
	trainAttrSpecs := base.ResolveAttributes(KNN.TrainingData, allNumericAttrs)

	// Reserve storage for most the most similar items
	distances := make(map[int]float64)

	// Reserve storage for voting map
	maxmapInt := make(map[string]int)
	maxmapFloat := make(map[string]float64)

	// Reserve storage for row computations
	trainRowBuf := make([]float64, len(allNumericAttrs))
	predRowBuf := make([]float64, len(allNumericAttrs))

	_, maxRow := what.Size()
	curRow := 0

	// build kdtree if algorithm is 'kdtree'
	kd := kdtree.New()
	srcRowNoMap := make([]int, 0)
	if KNN.Algorithm == "kdtree" {
		buildData := make([][]float64, 0)
		KNN.TrainingData.MapOverRows(trainAttrSpecs, func(trainRow [][]byte, srcRowNo int) (bool, error) {
			oneData := make([]float64, len(allNumericAttrs))
			// Read the float values out
			for i, _ := range allNumericAttrs {
				oneData[i] = base.UnpackBytesToFloat(trainRow[i])
			}
			srcRowNoMap = append(srcRowNoMap, srcRowNo)
			buildData = append(buildData, oneData)
			return true, nil
		})

		err := kd.Build(buildData)
		if err != nil {
			return nil, err
		}
	}

	// Iterate over all outer rows
	what.MapOverRows(whatAttrSpecs, func(predRow [][]byte, predRowNo int) (bool, error) {

		if (curRow%1) == 0 && curRow > 0 {
			fmt.Printf("KNN: %.2f %% done\r", float64(curRow)*100.0/float64(maxRow))
		}
		curRow++

		// Read the float values out
		for i, _ := range allNumericAttrs {
			predRowBuf[i] = base.UnpackBytesToFloat(predRow[i])
		}

		predMat := utilities.FloatsToMatrix(predRowBuf)

		switch KNN.Algorithm {
		case "linear":
			// Find the closest match in the training data
			KNN.TrainingData.MapOverRows(trainAttrSpecs, func(trainRow [][]byte, srcRowNo int) (bool, error) {
				// Read the float values out
				for i, _ := range allNumericAttrs {
					trainRowBuf[i] = base.UnpackBytesToFloat(trainRow[i])
				}

				// Compute the distance
				trainMat := utilities.FloatsToMatrix(trainRowBuf)
				distances[srcRowNo] = distanceFunc.Distance(predMat, trainMat)
				return true, nil
			})

			sorted := utilities.SortIntMap(distances)
			values := sorted[:KNN.NearestNeighbours]

			length := make([]float64, KNN.NearestNeighbours)
			for k, v := range values {
				length[k] = distances[v]
			}

			var maxClass string
			if KNN.Weighted {
				maxClass = KNN.weightedVote(maxmapFloat, values, length)
			} else {
				maxClass = KNN.vote(maxmapInt, values)
			}
			base.SetClass(ret, predRowNo, maxClass)

		case "kdtree":
			// search kdtree
			values, length, err := kd.Search(KNN.NearestNeighbours, distanceFunc, predRowBuf)
			if err != nil {
				return false, err
			}

			// map values to srcRowNo
			for k, v := range values {
				values[k] = srcRowNoMap[v]
			}

			var maxClass string
			if KNN.Weighted {
				maxClass = KNN.weightedVote(maxmapFloat, values, length)
			} else {
				maxClass = KNN.vote(maxmapInt, values)
			}
			base.SetClass(ret, predRowNo, maxClass)
		}

		return true, nil

	})

	return ret, nil
}

func (KNN *KNNClassifier) String() string {
	return fmt.Sprintf("KNNClassifier(%s, %d)", KNN.DistanceFunc, KNN.NearestNeighbours)
}

func (KNN *KNNClassifier) vote(maxmap map[string]int, values []int) string {
	// Reset maxMap
	for a := range maxmap {
		maxmap[a] = 0
	}

	// Refresh maxMap
	for _, elem := range values {
		label := base.GetClass(KNN.TrainingData, elem)
		if _, ok := maxmap[label]; ok {
			maxmap[label]++
		} else {
			maxmap[label] = 1
		}
	}

	// Sort the maxMap
	var maxClass string
	maxVal := -1
	for a := range maxmap {
		if maxmap[a] > maxVal {
			maxVal = maxmap[a]
			maxClass = a
		}
	}
	return maxClass
}

func (KNN *KNNClassifier) weightedVote(maxmap map[string]float64, values []int, length []float64) string {
	// Reset maxMap
	for a := range maxmap {
		maxmap[a] = 0
	}

	// Refresh maxMap
	for k, elem := range values {
		label := base.GetClass(KNN.TrainingData, elem)
		if _, ok := maxmap[label]; ok {
			maxmap[label] += (1 / length[k])
		} else {
			maxmap[label] = (1 / length[k])
		}
	}

	// Sort the maxMap
	var maxClass string
	maxVal := -1.0
	for a := range maxmap {
		if maxmap[a] > maxVal {
			maxVal = maxmap[a]
			maxClass = a
		}
	}
	return maxClass
}

// A KNNRegressor consists of a data matrix, associated result variables in the same order as the matrix, and a name.
type KNNRegressor struct {
	base.BaseEstimator
	Values       []float64
	DistanceFunc string
}

// NewKnnRegressor mints a new classifier.
func NewKnnRegressor(distfunc string) *KNNRegressor {
	KNN := KNNRegressor{}
	KNN.DistanceFunc = distfunc
	return &KNN
}

func (KNN *KNNRegressor) Fit(values []float64, numbers []float64, rows int, cols int) {
	if rows != len(values) {
		panic(matrix.ErrShape)
	}

	KNN.Data = mat64.NewDense(rows, cols, numbers)
	KNN.Values = values
}

func (KNN *KNNRegressor) Predict(vector *mat64.Dense, K int) float64 {
	// Get the number of rows
	rows, _ := KNN.Data.Dims()
	rownumbers := make(map[int]float64)
	labels := make([]float64, 0)

	// Check what distance function we are using
	var distanceFunc pairwise.PairwiseDistanceFunc
	switch KNN.DistanceFunc {
	case "euclidean":
		distanceFunc = pairwise.NewEuclidean()
	case "manhattan":
		distanceFunc = pairwise.NewManhattan()
	default:
		panic("unsupported distance function")
	}

	for i := 0; i < rows; i++ {
		row := KNN.Data.RowView(i)
		distance := distanceFunc.Distance(utilities.VectorToMatrix(row), vector)
		rownumbers[i] = distance
	}

	sorted := utilities.SortIntMap(rownumbers)
	values := sorted[:K]

	var sum float64
	for _, elem := range values {
		value := KNN.Values[elem]
		labels = append(labels, value)
		sum += value
	}

	average := sum / float64(K)
	return average
}
<|endoftext|>


<commit_before>package pachyderm

import (
	"context"
	"errors"
	"fmt"

	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/logical/framework"
	pclient "github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/auth"
)

// Revoke revokes the caller's credentials (by sending a request to Pachyderm).
// Unlike other handlers, it doesn't get assigned to a path; instead it's
// placed in Backend.Revoke in backend.go
func (b *backend) Revoke(ctx context.Context, req *logical.Request, data *framework.FieldData) (resp *logical.Response, retErr error) {
	b.Logger().Debug(fmt.Sprintf("(%s) %s received at %s", req.ID, req.Operation, req.Path))
	defer func() {
		b.Logger().Debug(fmt.Sprintf("(%s) %s finished at %s with result (success=%t)", req.ID, req.Operation, req.Path, retErr == nil && !resp.IsError()))
	}()

	// Extract pachyderm token from vault secret
	tokenIface, ok := req.Secret.InternalData["user_token"]
	if !ok {
		return nil, fmt.Errorf("secret is missing user_token")
	}
	userToken, ok := tokenIface.(string)
	if !ok {
		return nil, fmt.Errorf("secret.user_token has wrong type (expected string but was %T)", tokenIface)
	}

	// Get pach address and admin token from config
	config, err := getConfig(ctx, req.Storage)
	if err != nil {
		return nil, err
	}
	if len(config.AdminToken) == 0 {
		return nil, errors.New("plugin is missing admin token")
	}
	if len(config.PachdAddress) == 0 {
		return nil, errors.New("plugin is missing pachd address")
	}

	// Revoke creds
	err = revokeUserCredentials(ctx, config.PachdAddress, userToken, config.AdminToken)
	if err != nil {
		return nil, err
	}

	return &logical.Response{}, nil
}

// revokeUserCredentials revokes the Pachyderm authentication token 'userToken'
// using the vault plugin's Admin credentials.
func revokeUserCredentials(ctx context.Context, pachdAddress string, userToken string, adminToken string) error {
	// Setup a single use client w the given admin token / address
	client, err := pclient.NewFromAddress(pachdAddress)
	if err != nil {
		return err
	}
	client = client.WithCtx(ctx)
	client.SetAuthToken(adminToken)
	_, err = client.AuthAPIClient.RevokeAuthToken(client.Ctx(), &auth.RevokeAuthTokenRequest{
		Token: userToken,
	})
	return err
}
<commit_msg>Fix comment in revoke.go<commit_after>package pachyderm

import (
	"context"
	"errors"
	"fmt"

	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/logical/framework"
	pclient "github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/auth"
)

// Revoke revokes the caller's credentials (by sending a request to Pachyderm).
// Unlike other handlers, it doesn't get assigned to a path; instead it's
// called by the vault lease API when a token's lease expires or is revoked.
// It's set in Backend.Secrets[0].Revoke in backend.go
func (b *backend) Revoke(ctx context.Context, req *logical.Request, data *framework.FieldData) (resp *logical.Response, retErr error) {
	b.Logger().Debug(fmt.Sprintf("(%s) %s received at %s", req.ID, req.Operation, req.Path))
	defer func() {
		b.Logger().Debug(fmt.Sprintf("(%s) %s finished at %s with result (success=%t)", req.ID, req.Operation, req.Path, retErr == nil && !resp.IsError()))
	}()

	// Extract pachyderm token from vault secret
	tokenIface, ok := req.Secret.InternalData["user_token"]
	if !ok {
		return nil, fmt.Errorf("secret is missing user_token")
	}
	userToken, ok := tokenIface.(string)
	if !ok {
		return nil, fmt.Errorf("secret.user_token has wrong type (expected string but was %T)", tokenIface)
	}

	// Get pach address and admin token from config
	config, err := getConfig(ctx, req.Storage)
	if err != nil {
		return nil, err
	}
	if len(config.AdminToken) == 0 {
		return nil, errors.New("plugin is missing admin token")
	}
	if len(config.PachdAddress) == 0 {
		return nil, errors.New("plugin is missing pachd address")
	}

	// Revoke creds
	err = revokeUserCredentials(ctx, config.PachdAddress, userToken, config.AdminToken)
	if err != nil {
		return nil, err
	}

	return &logical.Response{}, nil
}

// revokeUserCredentials revokes the Pachyderm authentication token 'userToken'
// using the vault plugin's Admin credentials.
func revokeUserCredentials(ctx context.Context, pachdAddress string, userToken string, adminToken string) error {
	// Setup a single use client w the given admin token / address
	client, err := pclient.NewFromAddress(pachdAddress)
	if err != nil {
		return err
	}
	client = client.WithCtx(ctx)
	client.SetAuthToken(adminToken)
	_, err = client.AuthAPIClient.RevokeAuthToken(client.Ctx(), &auth.RevokeAuthTokenRequest{
		Token: userToken,
	})
	return err
}
<|endoftext|>


<commit_before>package libcentrifugo

import (
	"encoding/json"
	"errors"
	"github.com/shilkin/centrifugo/libcentrifugo/logger"
	"github.com/tarantool/go-tarantool"
	"strconv"
	"strings"
)

func (p *TarantoolPool) get() (conn *tarantool.Connection, err error) {
	if len(p.pool) == 0 {
		return nil, errors.New("Empty tarantool pool")
	}
	conn = p.pool[p.current]
	p.current++
	p.current = (p.current) % len(p.pool)
	return
}

type TarantoolEngine struct {
	app      *Application
	pool     *TarantoolPool
	endpoint string
}

type TarantoolEngineConfig struct {
	PoolConfig TarantoolPoolConfig
	Endpoint   string
}

type TarantoolPool struct {
	pool    []*tarantool.Connection
	config  TarantoolPoolConfig
	current int
}

type TarantoolPoolConfig struct {
	Address  string
	PoolSize int
	Opts     tarantool.Opts
}

/* MessageType
{
	"body": {
		"uid":"026c380d-13e1-47d9-42d2-e2dc0e41e8d5",
		"timestamp":"1440434259",
		"info":{
			"user":"3",
			"client":"83309b33-deb7-48ff-76c6-04b10e6a6523",
			"default_info":null,
			"channel_info": {
				"channel_extra_info_example":"you can add additional JSON data when authorizing"
			}
		},
		"channel":"$3_0",
		"data": {
				"Action":"mark",
				"Data":["00000000000000395684"]
			},
		"client":"83309b33-deb7-48ff-76c6-04b10e6a6523"
	},
	"error":null,
	"method":"message"
}
*/

type MessageType struct {
	Body   Message
	Error  string `json:error`
	Method string `json:method`
}

type ServiceMessage struct {
	Action string
	Data   []string
}

type IDs []string

func NewTarantoolEngine(app *Application, conf TarantoolEngineConfig) *TarantoolEngine {
	pool, err := newTarantoolPool(conf.PoolConfig)
	if err != nil {
		logger.FATAL.Fatalln(err)
	}

	e := &TarantoolEngine{
		app:      app,
		pool:     pool,
		endpoint: conf.Endpoint,
	}

	return e
}

func newTarantoolPool(config TarantoolPoolConfig) (p *TarantoolPool, err error) {
	if config.PoolSize == 0 {
		return nil, errors.New("Size of tarantool pool is zero")
	}

	p = &TarantoolPool{
		pool:   make([]*tarantool.Connection, config.PoolSize),
		config: config,
	}

	for i := 0; i < config.PoolSize; i++ {
		p.pool[i], err = tarantool.Connect(config.Address, config.Opts)
		if err != nil {
			return
		}
	}

	return p, nil
}

// getName returns a name of concrete engine implementation
func (e *TarantoolEngine) name() string {
	return "Tarantool"
}

// publish allows to send message into channel
func (e *TarantoolEngine) publish(chID ChannelID, message []byte) error {
	/*
		message:
			action: mark, push
			params:	[id,...]
	*/

	// Process service messages
	if chID != e.app.config.ControlChannel && chID != e.app.config.AdminChannel {
		if further, err := e.processMessage(chID, message); !further {
			return err // if no need further processing
		}
	}
	// All other messages
	return e.app.handleMsg(chID, message)
}

// subscribe on channel
func (e *TarantoolEngine) subscribe(chID ChannelID) (err error) {
	uid, ringno, project, err := parseChannelID(chID)
	if err != nil {
		return
	}

	conn, err := e.pool.get()
	if err != nil {
		logger.ERROR.Printf("subscribe tarantool pool error: %v\n", err.Error())
		return
	}

	_, err = conn.Call("notification_subscribe", []interface{}{uid, ringno, e.endpoint + "/api/" + project})

	return
}

// unsubscribe from channel
func (e *TarantoolEngine) unsubscribe(chID ChannelID) (err error) {
	uid, ringno, project, err := parseChannelID(chID)
	if err != nil {
		return
	}

	conn, err := e.pool.get()
	if err != nil {
		logger.ERROR.Printf("unsubscribe tarantool pool error: %v\n", err.Error())
		return
	}

	_, err = conn.Call("notification_unsubscribe", []interface{}{uid, ringno, e.endpoint + "/api/" + project})

	return
}

// addPresence sets or updates presence info for connection with uid
func (e *TarantoolEngine) addPresence(chID ChannelID, uid ConnID, info ClientInfo) (err error) {
	// not implemented
	return
}

// removePresence removes presence information for connection with uid
func (e *TarantoolEngine) removePresence(chID ChannelID, uid ConnID) (err error) {
	// not implemented
	return
}

// getPresence returns actual presence information for channel
func (e *TarantoolEngine) presence(chID ChannelID) (result map[ConnID]ClientInfo, err error) {
	// not implemented
	return
}

// addHistory adds message into channel history and takes care about history size
func (e *TarantoolEngine) addHistory(chID ChannelID, message Message, size, lifetime int64) (err error) {
	// not implemented
	return
}

// getHistory returns history messages for channel
// return empty slice
// all history pushed via publish
func (e *TarantoolEngine) history(chID ChannelID) (msgs []Message, err error) {
	uid, ringno, _, err := parseChannelID(chID)
	if err != nil {
		logger.ERROR.Printf("history parse chID error: %v\n", err.Error())
		return nil, err
	}

	conn, err := e.pool.get()
	if err != nil {
		logger.ERROR.Printf("history tarantool pool error: %v\n", err.Error())
		return nil, err
	}

	history, err := conn.Call("notification_read", []interface{}{uid, ringno})
	if err != nil {
		logger.ERROR.Printf("history notification_read error: %v\n", err.Error())
		return nil, err
	}

	return processHistory(history)
}

// helpers

type tarantoolHistoryItem struct {
	Count  interface{} `json:count`
	Status string      `json:status`
	ID     string      `json:id`
}

func processHistory(history *tarantool.Response) (msgs []Message, err error) {
	if len(history.Data) == 0 {
		return // history is empty
	}

	data := history.Data[0].([]interface{})
	if len(data) != 2 {
		return // history is empty
	}

	count := data[0]						// ring counter
	buffer := data[1].(string)				// string buffer
	ring := strings.Split(buffer[1:], ",")	// array of IDs

	if len(ring) == 0 {
		return	// history buffer is empty [useless?]
	}

	for _, id := range ring {
		encoded, err := json.Marshal(tarantoolHistoryItem{
			Count:  count, // redundancy in each item to pass number of unread notificatins
			Status: string(id[0]),
			ID:     string(id[1:]),
		})
		if err != nil {
			logger.ERROR.Println(err)
			continue
		}
		rawMessage := json.RawMessage([]byte(encoded))
		msgs = append(msgs, Message{Data: &rawMessage})
	}

	return
}

func (e *TarantoolEngine) processMessage(chID ChannelID, message []byte) (needFurtherProcessing bool, err error) {
	var msg MessageType
	err = json.Unmarshal(message, &msg)
	if err != nil {
		return true, err
	}

	var srv ServiceMessage
	err = json.Unmarshal(*msg.Body.Data, &srv)
	if err != nil {
		return true, err
	}

	if srv.Action == "" {
		return true, nil
	}

	var functionName string
	switch srv.Action {
	case "mark":
		functionName = "notification_mark"
	case "push":
		functionName = "notification_push"
	default:
		return true, nil
	}

	var uid, ringno int64
	uid, ringno, _, err = parseChannelID(chID)
	if err != nil {
		return
	}

	var conn *tarantool.Connection
	conn, err = e.pool.get()
	if err != nil {
		return
	}

	for _, id := range srv.Data {
		_, err = conn.Call(functionName, []interface{}{uid, ringno, id})
		if err != nil {
			logger.ERROR.Printf("%s call error: %s", functionName, err)
			return
		}
	}

	return
}

func parseChannelID(chID ChannelID) (uid, ringno int64, project string, err error) {
	// split chID <centrifugo>.<project>.[$]<uid>_<ringno>
	str := string(chID)
	logger.DEBUG.Printf("parseChannelID %s", str)

	result := strings.Split(str, ".")
	if len(result) != 3 {
		logger.DEBUG.Printf("unexpected ChannelID %s", str)
		return
	}

	project = result[1]
	str = result[2]

	separator := "_"
	prefix := "$"

	if strings.HasPrefix(str, prefix) {
		str = strings.TrimLeft(str, prefix)
	}
	channel := strings.Split(str, separator)

	uid, err = strconv.ParseInt(channel[0], 10, 64)
	if err != nil {
		return
	}
	ringno, err = strconv.ParseInt(channel[1], 10, 64)
	if err != nil {
		return
	}
	return
}
<commit_msg>Additional logging during tarantool connection pool initialization<commit_after>package libcentrifugo

import (
	"encoding/json"
	"errors"
	"github.com/shilkin/centrifugo/libcentrifugo/logger"
	"github.com/tarantool/go-tarantool"
	"strconv"
	"strings"
)

func (p *TarantoolPool) get() (conn *tarantool.Connection, err error) {
	if len(p.pool) == 0 {
		return nil, errors.New("Empty tarantool pool")
	}
	conn = p.pool[p.current]
	p.current++
	p.current = (p.current) % len(p.pool)
	return
}

type TarantoolEngine struct {
	app      *Application
	pool     *TarantoolPool
	endpoint string
}

type TarantoolEngineConfig struct {
	PoolConfig TarantoolPoolConfig
	Endpoint   string
}

type TarantoolPool struct {
	pool    []*tarantool.Connection
	config  TarantoolPoolConfig
	current int
}

type TarantoolPoolConfig struct {
	Address  string
	PoolSize int
	Opts     tarantool.Opts
}

/* MessageType
{
	"body": {
		"uid":"026c380d-13e1-47d9-42d2-e2dc0e41e8d5",
		"timestamp":"1440434259",
		"info":{
			"user":"3",
			"client":"83309b33-deb7-48ff-76c6-04b10e6a6523",
			"default_info":null,
			"channel_info": {
				"channel_extra_info_example":"you can add additional JSON data when authorizing"
			}
		},
		"channel":"$3_0",
		"data": {
				"Action":"mark",
				"Data":["00000000000000395684"]
			},
		"client":"83309b33-deb7-48ff-76c6-04b10e6a6523"
	},
	"error":null,
	"method":"message"
}
*/

type MessageType struct {
	Body   Message
	Error  string `json:error`
	Method string `json:method`
}

type ServiceMessage struct {
	Action string
	Data   []string
}

type IDs []string

func NewTarantoolEngine(app *Application, conf TarantoolEngineConfig) *TarantoolEngine {
	logger.INFO.Printf("Initializing tarantool connection pool...")
	pool, err := newTarantoolPool(conf.PoolConfig)
	if err != nil {
		logger.FATAL.Fatalln(err)
	}

	e := &TarantoolEngine{
		app:      app,
		pool:     pool,
		endpoint: conf.Endpoint,
	}

	return e
}

func newTarantoolPool(config TarantoolPoolConfig) (p *TarantoolPool, err error) {
	if config.PoolSize == 0 {
		return nil, errors.New("Size of tarantool pool is zero")
	}

	p = &TarantoolPool{
		pool:   make([]*tarantool.Connection, config.PoolSize),
		config: config,
	}

	for i := 0; i < config.PoolSize; i++ {
		logger.INFO.Printf("[%d] Connecting to tarantool on %s...", i, config.Address)
		p.pool[i], err = tarantool.Connect(config.Address, config.Opts)
		if err != nil {
			return
		}
		logger.INFO.Printf("[%d] Connected to tarantool on %s", i, config.Address)
	}

	return p, nil
}

// getName returns a name of concrete engine implementation
func (e *TarantoolEngine) name() string {
	return "Tarantool"
}

// publish allows to send message into channel
func (e *TarantoolEngine) publish(chID ChannelID, message []byte) error {
	/*
		message:
			action: mark, push
			params:	[id,...]
	*/

	// Process service messages
	if chID != e.app.config.ControlChannel && chID != e.app.config.AdminChannel {
		if further, err := e.processMessage(chID, message); !further {
			return err // if no need further processing
		}
	}
	// All other messages
	return e.app.handleMsg(chID, message)
}

// subscribe on channel
func (e *TarantoolEngine) subscribe(chID ChannelID) (err error) {
	uid, ringno, project, err := parseChannelID(chID)
	if err != nil {
		return
	}

	conn, err := e.pool.get()
	if err != nil {
		logger.ERROR.Printf("subscribe tarantool pool error: %v\n", err.Error())
		return
	}

	_, err = conn.Call("notification_subscribe", []interface{}{uid, ringno, e.endpoint + "/api/" + project})

	return
}

// unsubscribe from channel
func (e *TarantoolEngine) unsubscribe(chID ChannelID) (err error) {
	uid, ringno, project, err := parseChannelID(chID)
	if err != nil {
		return
	}

	conn, err := e.pool.get()
	if err != nil {
		logger.ERROR.Printf("unsubscribe tarantool pool error: %v\n", err.Error())
		return
	}

	_, err = conn.Call("notification_unsubscribe", []interface{}{uid, ringno, e.endpoint + "/api/" + project})

	return
}

// addPresence sets or updates presence info for connection with uid
func (e *TarantoolEngine) addPresence(chID ChannelID, uid ConnID, info ClientInfo) (err error) {
	// not implemented
	return
}

// removePresence removes presence information for connection with uid
func (e *TarantoolEngine) removePresence(chID ChannelID, uid ConnID) (err error) {
	// not implemented
	return
}

// getPresence returns actual presence information for channel
func (e *TarantoolEngine) presence(chID ChannelID) (result map[ConnID]ClientInfo, err error) {
	// not implemented
	return
}

// addHistory adds message into channel history and takes care about history size
func (e *TarantoolEngine) addHistory(chID ChannelID, message Message, size, lifetime int64) (err error) {
	// not implemented
	return
}

// getHistory returns history messages for channel
// return empty slice
// all history pushed via publish
func (e *TarantoolEngine) history(chID ChannelID) (msgs []Message, err error) {
	uid, ringno, _, err := parseChannelID(chID)
	if err != nil {
		logger.ERROR.Printf("history parse chID error: %v\n", err.Error())
		return nil, err
	}

	conn, err := e.pool.get()
	if err != nil {
		logger.ERROR.Printf("history tarantool pool error: %v\n", err.Error())
		return nil, err
	}

	history, err := conn.Call("notification_read", []interface{}{uid, ringno})
	if err != nil {
		logger.ERROR.Printf("history notification_read error: %v\n", err.Error())
		return nil, err
	}

	return processHistory(history)
}

// helpers

type tarantoolHistoryItem struct {
	Count  interface{} `json:count`
	Status string      `json:status`
	ID     string      `json:id`
}

func processHistory(history *tarantool.Response) (msgs []Message, err error) {
	if len(history.Data) == 0 {
		return // history is empty
	}

	data := history.Data[0].([]interface{})
	if len(data) != 2 {
		return // history is empty
	}

	count := data[0]						// ring counter
	buffer := data[1].(string)				// string buffer
	ring := strings.Split(buffer[1:], ",")	// array of IDs

	if len(ring) == 0 {
		return	// history buffer is empty [useless?]
	}

	for _, id := range ring {
		encoded, err := json.Marshal(tarantoolHistoryItem{
			Count:  count, // redundancy in each item to pass number of unread notificatins
			Status: string(id[0]),
			ID:     string(id[1:]),
		})
		if err != nil {
			logger.ERROR.Println(err)
			continue
		}
		rawMessage := json.RawMessage([]byte(encoded))
		msgs = append(msgs, Message{Data: &rawMessage})
	}

	return
}

func (e *TarantoolEngine) processMessage(chID ChannelID, message []byte) (needFurtherProcessing bool, err error) {
	var msg MessageType
	err = json.Unmarshal(message, &msg)
	if err != nil {
		return true, err
	}

	var srv ServiceMessage
	err = json.Unmarshal(*msg.Body.Data, &srv)
	if err != nil {
		return true, err
	}

	if srv.Action == "" {
		return true, nil
	}

	var functionName string
	switch srv.Action {
	case "mark":
		functionName = "notification_mark"
	case "push":
		functionName = "notification_push"
	default:
		return true, nil
	}

	var uid, ringno int64
	uid, ringno, _, err = parseChannelID(chID)
	if err != nil {
		return
	}

	var conn *tarantool.Connection
	conn, err = e.pool.get()
	if err != nil {
		return
	}

	for _, id := range srv.Data {
		_, err = conn.Call(functionName, []interface{}{uid, ringno, id})
		if err != nil {
			logger.ERROR.Printf("%s call error: %s", functionName, err)
			return
		}
	}

	return
}

func parseChannelID(chID ChannelID) (uid, ringno int64, project string, err error) {
	// split chID <centrifugo>.<project>.[$]<uid>_<ringno>
	str := string(chID)
	logger.DEBUG.Printf("parseChannelID %s", str)

	result := strings.Split(str, ".")
	if len(result) != 3 {
		logger.DEBUG.Printf("unexpected ChannelID %s", str)
		return
	}

	project = result[1]
	str = result[2]

	separator := "_"
	prefix := "$"

	if strings.HasPrefix(str, prefix) {
		str = strings.TrimLeft(str, prefix)
	}
	channel := strings.Split(str, separator)

	uid, err = strconv.ParseInt(channel[0], 10, 64)
	if err != nil {
		return
	}
	ringno, err = strconv.ParseInt(channel[1], 10, 64)
	if err != nil {
		return
	}
	return
}
<|endoftext|>


<commit_before>package mbr

// Type constants for the GUID for type of partition, see https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_entries
type Type byte

// List of GUID partition types
const (
	Empty         Type = 0x00
	Fat12         Type = 0x01
	XenixRoot     Type = 0x02
	XenixUsr      Type = 0x03
	Fat16         Type = 0x04
	ExtendedCHS   Type = 0x05
	Fat16b        Type = 0x06
	NTFS          Type = 0x07
	CommodoreFAT  Type = 0x08
	Fat32CHS      Type = 0x0b
	Fat32LBA      Type = 0x0c
	Fat16bLBA     Type = 0x0e
	ExtendedLBA   Type = 0x0f
	Linux         Type = 0x83
	LinuxExtended Type = 0x85
	LinuxLVM      Type = 0x8e
	Iso9660       Type = 0x96
	MacOSXUFS     Type = 0xa8
	MacOSXBoot    Type = 0xab
	HFS           Type = 0xaf
	Solaris8Boot  Type = 0xbe
	GPTProtective Type = 0xee
	EFISystem     Type = 0xef
	VMWareFS      Type = 0xfb
	VMWareSwap    Type = 0xfc
)
<commit_msg>add linux swap partition type (#155)<commit_after>package mbr

// Type constants for the GUID for type of partition, see https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_entries
type Type byte

// List of GUID partition types
const (
	Empty         Type = 0x00
	Fat12         Type = 0x01
	XenixRoot     Type = 0x02
	XenixUsr      Type = 0x03
	Fat16         Type = 0x04
	ExtendedCHS   Type = 0x05
	Fat16b        Type = 0x06
	NTFS          Type = 0x07
	CommodoreFAT  Type = 0x08
	Fat32CHS      Type = 0x0b
	Fat32LBA      Type = 0x0c
	Fat16bLBA     Type = 0x0e
	ExtendedLBA   Type = 0x0f
	LinuxSwap     Type = 0x82
	Linux         Type = 0x83
	LinuxExtended Type = 0x85
	LinuxLVM      Type = 0x8e
	Iso9660       Type = 0x96
	MacOSXUFS     Type = 0xa8
	MacOSXBoot    Type = 0xab
	HFS           Type = 0xaf
	Solaris8Boot  Type = 0xbe
	GPTProtective Type = 0xee
	EFISystem     Type = 0xef
	VMWareFS      Type = 0xfb
	VMWareSwap    Type = 0xfc
)
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"github.com/gin-gonic/gin"
	"github.com/jmhodges/levigo"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
)

var (
	serveaddr = flag.String(
		"serveaddr",
		"127.0.0.1:8070",
		"[host]:port or /path/to/socketfile of where to run the server",
	)
)

var (
	db *levigo.DB
	ro *levigo.ReadOptions
	wo *levigo.WriteOptions
)

func main() {
	flag.Parse()

	openDB()
	defer wo.Close()
	defer ro.Close()
	defer db.Close()

	engine := initEngine()
	log.Print(run(engine))
}

func initEngine() *gin.Engine {
	engine := gin.New()
	engine.Use(gin.Recovery())

	engine.GET("/key/:name", func(c *gin.Context) {
		b, err := db.Get(ro, []byte(c.Params.ByName("name")))
		if err != nil {
			c.Fail(500, err)
		} else if b == nil {
			c.AbortWithStatus(404)
		} else {
			c.String(200, string(b))
		}
	})

	engine.PUT("/key/:name", func(c *gin.Context) {
		buf := &bytes.Buffer{}
		if _, err := io.Copy(buf, c.Request.Body); err != nil {
			c.Fail(500, err)
			return
		}

		err := db.Put(wo, []byte(c.Params.ByName("name")), buf.Bytes())
		if err != nil {
			c.Fail(500, err)
		} else {
			c.Writer.WriteHeader(204)
		}
	})

	engine.POST("/snapshot", func(c *gin.Context) {
		req := &struct {
			Destination string
		}{}
		if !c.Bind(req) {
			return
		}

		if err := makeSnap(req.Destination); err != nil {
			c.Fail(500, err)
		} else {
			c.Writer.WriteHeader(204)
		}
	})

	return engine
}

func run(engine *gin.Engine) error {
	if strings.Contains(*serveaddr, ":") {
		return engine.Run(*serveaddr)
	}

	listener, err := net.Listen("unix", *serveaddr)
	if err != nil {
		log.Fatal(err)
	}

	return (&http.Server{Handler: engine}).Serve(listener)
}

func openDB() {
	args := flag.Args()
	if len(args) == 0 {
		log.Fatal("missing db path cmdline argument")
	}
	path := args[0]

	opts := levigo.NewOptions()
	opts.SetCreateIfMissing(true)
	ldb, err := levigo.Open(path, opts)
	if err != nil {
		log.Fatalf("opening leveldb: %s", err)
	}

	db = ldb
	ro = levigo.NewReadOptions()
	wo = levigo.NewWriteOptions()
}

func makeSnap(dest string) error {
	opts := levigo.NewOptions()
	opts.SetCreateIfMissing(true)
	opts.SetErrorIfExists(true)
	to, err := levigo.Open(dest, opts)
	if err != nil {
		return err
	}
	defer to.Close()

	ss := db.NewSnapshot()
	ro := levigo.NewReadOptions()
	ro.SetSnapshot(ss)
	ro.SetFillCache(false)

	it := db.NewIterator(ro)
	defer it.Close()

	wb := levigo.NewWriteBatch()

	var i uint
	for it.SeekToFirst(); it.Valid(); it.Next() {
		wb.Put(it.Key(), it.Value())
		i++

		if i%1000 == 0 {
			wb, err = dumpBatch(wb, to, true)
			if err != nil {
				goto fail
			}
		}
	}

	if i%1000 != 0 {
		_, err = dumpBatch(wb, to, false)
		if err != nil {
			goto fail
		}
	}

	return nil

fail:
	levigo.DestroyDatabase(dest, opts)
	return err
}

func dumpBatch(wb *levigo.WriteBatch, dest *levigo.DB, more bool) (*levigo.WriteBatch, error) {
	defer wb.Close()

	err := dest.Write(wo, wb)
	if err != nil {
		return nil, err
	}

	if more {
		return levigo.NewWriteBatch(), nil
	}
	return nil, nil
}
<commit_msg>drop gin and only depend on raw httprouter<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"github.com/julienschmidt/httprouter"
	"github.com/jmhodges/levigo"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
)

var (
	serveaddr = flag.String(
		"serveaddr",
		"127.0.0.1:8070",
		"[host]:port or /path/to/socketfile of where to run the server",
	)
)

var (
	db *levigo.DB
	ro *levigo.ReadOptions
	wo *levigo.WriteOptions
)

func main() {
	flag.Parse()

	openDB()
	defer wo.Close()
	defer ro.Close()
	defer db.Close()

	router := initRouter()
	log.Print(run(router))
}

func initRouter() *httprouter.Router {
	router := &httprouter.Router{
		// precision in urls -- I'd rather know when my client is wrong
		RedirectTrailingSlash: false,
		RedirectFixedPath: false,

		HandleMethodNotAllowed: true,
		PanicHandler: handlePanics,
	}

	router.GET("/key/:name", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		b, err := db.Get(ro, []byte(p.ByName("name")))
		if err != nil {
			failErr(w, err)
		} else if b == nil {
			failCode(w, http.StatusNotFound)
			w.WriteHeader(404)
		} else {
			w.Write(b)
		}
	})

	router.PUT("/key/:name", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		buf := &bytes.Buffer{}
		if _, err := io.Copy(buf, r.Body); err != nil {
			failErr(w, err)
			return
		}

		err := db.Put(wo, []byte(p.ByName("name")), buf.Bytes())
		if err != nil {
			failErr(w, err)
		} else {
			w.WriteHeader(204)
		}
	})

	router.POST("/snapshot", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		req := &struct {
			Destination string
		}{}
		err := json.NewDecoder(r.Body).Decode(req)
		if err != nil {
			failErr(w, err)
			return
		}

		if err := makeSnap(req.Destination); err != nil {
			failErr(w, err)
		} else {
			w.WriteHeader(204)
		}
	})

	return router
}

func run(router *httprouter.Router) error {
	if strings.Contains(*serveaddr, ":") {
		return http.ListenAndServe(*serveaddr, router)
	}

	listener, err := net.Listen("unix", *serveaddr)
	if err != nil {
		log.Fatal(err)
	}

	return (&http.Server{Handler: router}).Serve(listener)
}

func openDB() {
	args := flag.Args()
	if len(args) == 0 {
		log.Fatal("missing db path cmdline argument")
	}
	path := args[0]

	opts := levigo.NewOptions()
	opts.SetCreateIfMissing(true)
	ldb, err := levigo.Open(path, opts)
	if err != nil {
		log.Fatalf("opening leveldb: %s", err)
	}

	db = ldb
	ro = levigo.NewReadOptions()
	wo = levigo.NewWriteOptions()
}

func makeSnap(dest string) error {
	opts := levigo.NewOptions()
	opts.SetCreateIfMissing(true)
	opts.SetErrorIfExists(true)
	to, err := levigo.Open(dest, opts)
	if err != nil {
		return err
	}
	defer to.Close()

	ss := db.NewSnapshot()
	ro := levigo.NewReadOptions()
	ro.SetSnapshot(ss)
	ro.SetFillCache(false)

	it := db.NewIterator(ro)
	defer it.Close()

	wb := levigo.NewWriteBatch()

	var i uint
	for it.SeekToFirst(); it.Valid(); it.Next() {
		wb.Put(it.Key(), it.Value())
		i++

		if i%1000 == 0 {
			wb, err = dumpBatch(wb, to, true)
			if err != nil {
				goto fail
			}
		}
	}

	if i%1000 != 0 {
		_, err = dumpBatch(wb, to, false)
		if err != nil {
			goto fail
		}
	}

	return nil

fail:
	levigo.DestroyDatabase(dest, opts)
	return err
}

func dumpBatch(wb *levigo.WriteBatch, dest *levigo.DB, more bool) (*levigo.WriteBatch, error) {
	defer wb.Close()

	err := dest.Write(wo, wb)
	if err != nil {
		return nil, err
	}

	if more {
		return levigo.NewWriteBatch(), nil
	}
	return nil, nil
}

func handlePanics(w http.ResponseWriter, r *http.Request, err interface{}) {
	log.Printf("PANIC in handler: %s", err)
	w.WriteHeader(http.StatusInternalServerError)
}

func failErr(w http.ResponseWriter, err error) {
	log.Print(err)
	w.WriteHeader(http.StatusInternalServerError)
}

func failCode(w http.ResponseWriter, code int) {
	w.WriteHeader(code)
	w.Write([]byte(http.StatusText(code)))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/ark-lang/ark/src/util/log"

	"gopkg.in/alecthomas/kingpin.v2"

	"github.com/ark-lang/ark/src/codegen"
	"github.com/ark-lang/ark/src/codegen/LLVMCodegen"
	"github.com/ark-lang/ark/src/doc"
	"github.com/ark-lang/ark/src/lexer"
	"github.com/ark-lang/ark/src/parser"
	"github.com/ark-lang/ark/src/util"
)

const (
	VERSION = "0.0.2"
	AUTHOR  = "The Ark Authors"
)

var startTime time.Time

func main() {
	startTime = time.Now()

	command := kingpin.MustParse(app.Parse(os.Args[1:]))
	log.SetLevel(*logLevel)
	log.SetTags(*logTags)

	switch command {
	case buildCom.FullCommand():
		if len(*buildInputs) == 0 {
			setupErr("No input files passed.")
		}

		ccArgs := []string{}
		if *buildStatic {
			ccArgs = append(ccArgs, "-static")
		}

		outputType := parseOutputType(*buildOutputType)
		build(*buildInputs, *buildOutput, *buildCodegen, ccArgs, outputType)
		printFinishedMessage(startTime, buildCom.FullCommand(), len(*buildInputs))
		if *buildRun {
			if outputType != LLVMCodegen.OUTPUT_EXECUTABLE {
				setupErr("Can only use --run flag when building executable")
			}
			run(*buildOutput)
		}

	case docgenCom.FullCommand():
		docgen(*docgenInputs, *docgenDir)
		printFinishedMessage(startTime, docgenCom.FullCommand(), len(*docgenInputs))
	}
}

func printFinishedMessage(startTime time.Time, command string, numFiles int) {
	dur := time.Since(startTime)
	log.Info("main", "%s (%d file(s), %.2fms)\n",
		util.TEXT_GREEN+util.TEXT_BOLD+fmt.Sprintf("Finished %s", command)+util.TEXT_RESET,
		numFiles, float32(dur.Nanoseconds())/1000000)
}

func setupErr(err string, stuff ...interface{}) {
	log.Error("main", util.TEXT_RED+util.TEXT_BOLD+"Setup error:"+util.TEXT_RESET+" %s\n",
		fmt.Sprintf(err, stuff...))
	os.Exit(util.EXIT_FAILURE_SETUP)
}

func build(files []string, outputFile string, cg string, ccArgs []string, outputType LLVMCodegen.OutputType) {
	// read source files
	var sourcefiles []*lexer.Sourcefile

	timed("reading sourcefiles", func() {
		for _, file := range files {
			sourcefile, err := lexer.NewSourcefile(file)
			if err != nil {
				setupErr("%s", err.Error())
			}
			sourcefiles = append(sourcefiles, sourcefile)
		}
	})

	// lexing
	timed("lexing phase", func() {
		for _, file := range sourcefiles {
			file.Tokens = lexer.Lex(file)
		}
	})

	// parsing
	var parsedFiles []*parser.ParseTree
	timed("parsing phase", func() {
		for _, file := range sourcefiles {
			parsedFiles = append(parsedFiles, parser.Parse(file))
		}
	})

	// construction
	var constructedModules []*parser.Module
	modules := make(map[string]*parser.Module)
	timed("construction phase", func() {
		for _, file := range parsedFiles {
			constructedModules = append(constructedModules, parser.Construct(file, modules))
		}
	})

	// resolve
	timed("resolve phase", func() {
		// TODO: We're looping over a map, the order we get is thus random
		for _, module := range modules {
			res := &parser.Resolver{Module: module}
			res.Resolve(modules)
		}
	})

	// semantic analysis
	timed("semantic analysis phase", func() {
		// TODO: We're looping over a map, the order we get is thus random
		for _, module := range modules {
			sem := &parser.SemanticAnalyzer{Module: module}
			sem.Analyze(modules)
		}
	})

	// codegen
	if cg != "none" {
		var gen codegen.Codegen

		switch cg {
		case "llvm":
			gen = &LLVMCodegen.Codegen{
				OutputName:   outputFile,
				CompilerArgs: ccArgs,
				OutputType:   outputType,
			}
		default:
			log.Error("main", util.Red("error: ")+"Invalid backend choice `"+cg+"`")
			os.Exit(1)
		}

		timed("codegen phase", func() {
			gen.Generate(constructedModules, modules)
		})
	}

}

func timed(title string, fn func()) {
	log.Verboseln("main", util.TEXT_BOLD+util.TEXT_GREEN+"Started "+title+util.TEXT_RESET)
	start := time.Now()

	fn()

	duration := time.Since(start)
	log.Verboseln("main", util.TEXT_BOLD+util.TEXT_GREEN+"Ended "+title+util.TEXT_RESET+" (%.2fms)", float32(duration)/1000000)
}

func run(output string) {
	cmd := exec.Command("./" + output)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
}

func docgen(input []string, dir string) {
	// TODO: eww duplication
	// read source files
	var sourcefiles []*lexer.Sourcefile

	timed("reading sourcefiles", func() {
		for _, file := range input {
			sourcefile, err := lexer.NewSourcefile(file)
			if err != nil {
				setupErr("%s", err.Error())
			}
			sourcefiles = append(sourcefiles, sourcefile)
		}
	})

	// lexing
	timed("lexing phase", func() {
		for _, file := range sourcefiles {
			file.Tokens = lexer.Lex(file)
		}
	})

	// parsing
	var parsedFiles []*parser.ParseTree

	timed("parsing phase", func() {
		for _, file := range sourcefiles {
			parsedFiles = append(parsedFiles, parser.Parse(file))
		}
	})

	gen := &doc.Docgen{
		Input: parsedFiles,
		Dir:   dir,
	}

	gen.Generate()
}
<commit_msg>Deduplicate parsing code in main.go<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/ark-lang/ark/src/util/log"

	"gopkg.in/alecthomas/kingpin.v2"

	"github.com/ark-lang/ark/src/codegen"
	"github.com/ark-lang/ark/src/codegen/LLVMCodegen"
	"github.com/ark-lang/ark/src/doc"
	"github.com/ark-lang/ark/src/lexer"
	"github.com/ark-lang/ark/src/parser"
	"github.com/ark-lang/ark/src/util"
)

const (
	VERSION = "0.0.2"
	AUTHOR  = "The Ark Authors"
)

var startTime time.Time

func main() {
	startTime = time.Now()

	command := kingpin.MustParse(app.Parse(os.Args[1:]))
	log.SetLevel(*logLevel)
	log.SetTags(*logTags)

	switch command {
	case buildCom.FullCommand():
		if len(*buildInputs) == 0 {
			setupErr("No input files passed.")
		}

		ccArgs := []string{}
		if *buildStatic {
			ccArgs = append(ccArgs, "-static")
		}

		outputType := parseOutputType(*buildOutputType)
		build(*buildInputs, *buildOutput, *buildCodegen, ccArgs, outputType)
		printFinishedMessage(startTime, buildCom.FullCommand(), len(*buildInputs))
		if *buildRun {
			if outputType != LLVMCodegen.OUTPUT_EXECUTABLE {
				setupErr("Can only use --run flag when building executable")
			}
			run(*buildOutput)
		}

	case docgenCom.FullCommand():
		docgen(*docgenInputs, *docgenDir)
		printFinishedMessage(startTime, docgenCom.FullCommand(), len(*docgenInputs))
	}
}

func printFinishedMessage(startTime time.Time, command string, numFiles int) {
	dur := time.Since(startTime)
	log.Info("main", "%s (%d file(s), %.2fms)\n",
		util.TEXT_GREEN+util.TEXT_BOLD+fmt.Sprintf("Finished %s", command)+util.TEXT_RESET,
		numFiles, float32(dur.Nanoseconds())/1000000)
}

func setupErr(err string, stuff ...interface{}) {
	log.Error("main", util.TEXT_RED+util.TEXT_BOLD+"Setup error:"+util.TEXT_RESET+" %s\n",
		fmt.Sprintf(err, stuff...))
	os.Exit(util.EXIT_FAILURE_SETUP)
}

func parseFiles(files []string) ([]*parser.Module, map[string]*parser.Module) {
	// read source files
	var sourcefiles []*lexer.Sourcefile

	timed("reading sourcefiles", func() {
		for _, file := range files {
			sourcefile, err := lexer.NewSourcefile(file)
			if err != nil {
				setupErr("%s", err.Error())
			}
			sourcefiles = append(sourcefiles, sourcefile)
		}
	})

	// lexing
	timed("lexing phase", func() {
		for _, file := range sourcefiles {
			file.Tokens = lexer.Lex(file)
		}
	})

	// parsing
	var parsedFiles []*parser.ParseTree
	timed("parsing phase", func() {
		for _, file := range sourcefiles {
			parsedFiles = append(parsedFiles, parser.Parse(file))
		}
	})

	// construction
	var constructedModules []*parser.Module
	modules := make(map[string]*parser.Module)
	timed("construction phase", func() {
		for _, file := range parsedFiles {
			constructedModules = append(constructedModules, parser.Construct(file, modules))
		}
	})

	return constructedModules, modules
}

func build(files []string, outputFile string, cg string, ccArgs []string, outputType LLVMCodegen.OutputType) {
	constructedModules, modules := parseFiles(files)

	// resolve
	timed("resolve phase", func() {
		// TODO: We're looping over a map, the order we get is thus random
		for _, module := range modules {
			res := &parser.Resolver{Module: module}
			res.Resolve(modules)
		}
	})

	// semantic analysis
	timed("semantic analysis phase", func() {
		// TODO: We're looping over a map, the order we get is thus random
		for _, module := range modules {
			sem := &parser.SemanticAnalyzer{Module: module}
			sem.Analyze(modules)
		}
	})

	// codegen
	if cg != "none" {
		var gen codegen.Codegen

		switch cg {
		case "llvm":
			gen = &LLVMCodegen.Codegen{
				OutputName:   outputFile,
				CompilerArgs: ccArgs,
				OutputType:   outputType,
			}
		default:
			log.Error("main", util.Red("error: ")+"Invalid backend choice `"+cg+"`")
			os.Exit(1)
		}

		timed("codegen phase", func() {
			gen.Generate(constructedModules, modules)
		})
	}

}

func timed(title string, fn func()) {
	log.Verboseln("main", util.TEXT_BOLD+util.TEXT_GREEN+"Started "+title+util.TEXT_RESET)
	start := time.Now()

	fn()

	duration := time.Since(start)
	log.Verboseln("main", util.TEXT_BOLD+util.TEXT_GREEN+"Ended "+title+util.TEXT_RESET+" (%.2fms)", float32(duration)/1000000)
}

func run(output string) {
	cmd := exec.Command("./" + output)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
}

func docgen(input []string, dir string) {
	constructedModules, _ := parseFiles(input)

	gen := &doc.Docgen{
		Input: constructedModules,
		Dir:   dir,
	}

	gen.Generate()
}
<|endoftext|>


<commit_before>// +build linux

package fs2

import (
	"path/filepath"
	"strings"

	"github.com/opencontainers/runc/libcontainer/cgroups"
	"github.com/opencontainers/runc/libcontainer/cgroups/fscommon"
	"github.com/opencontainers/runc/libcontainer/configs"
	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

func isPidsSet(cgroup *configs.Cgroup) bool {
	return cgroup.Resources.PidsLimit != 0
}

func setPids(dirPath string, cgroup *configs.Cgroup) error {
	if !isPidsSet(cgroup) {
		return nil
	}
	if val := numToStr(cgroup.Resources.PidsLimit); val != "" {
		if err := fscommon.WriteFile(dirPath, "pids.max", val); err != nil {
			return err
		}
	}

	return nil
}

func statPidsWithoutController(dirPath string, stats *cgroups.Stats) error {
	// if the controller is not enabled, let's read PIDS from cgroups.procs
	// (or threads if cgroup.threads is enabled)
	contents, err := fscommon.ReadFile(dirPath, "cgroup.procs")
	if errors.Is(err, unix.ENOTSUP) {
		contents, err = fscommon.ReadFile(dirPath, "cgroup.threads")
	}
	if err != nil {
		return err
	}
	pids := make(map[string]string)
	for _, i := range strings.Split(contents, "\n") {
		if i != "" {
			pids[i] = i
		}
	}
	stats.PidsStats.Current = uint64(len(pids))
	stats.PidsStats.Limit = 0
	return nil
}

func statPids(dirPath string, stats *cgroups.Stats) error {
	current, err := fscommon.GetCgroupParamUint(dirPath, "pids.current")
	if err != nil {
		return errors.Wrap(err, "failed to parse pids.current")
	}

	maxString, err := fscommon.GetCgroupParamString(dirPath, "pids.max")
	if err != nil {
		return errors.Wrap(err, "failed to parse pids.max")
	}

	// Default if pids.max == "max" is 0 -- which represents "no limit".
	var max uint64
	if maxString != "max" {
		max, err = fscommon.ParseUint(maxString, 10, 64)
		if err != nil {
			return errors.Wrapf(err, "failed to parse pids.max - unable to parse %q as a uint from Cgroup file %q",
				maxString, filepath.Join(dirPath, "pids.max"))
		}
	}

	stats.PidsStats.Current = current
	stats.PidsStats.Limit = max
	return nil
}
<commit_msg>libct/cg/fs2/getPidsWithoutController: optimize<commit_after>// +build linux

package fs2

import (
	"path/filepath"
	"strings"

	"github.com/opencontainers/runc/libcontainer/cgroups"
	"github.com/opencontainers/runc/libcontainer/cgroups/fscommon"
	"github.com/opencontainers/runc/libcontainer/configs"
	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

func isPidsSet(cgroup *configs.Cgroup) bool {
	return cgroup.Resources.PidsLimit != 0
}

func setPids(dirPath string, cgroup *configs.Cgroup) error {
	if !isPidsSet(cgroup) {
		return nil
	}
	if val := numToStr(cgroup.Resources.PidsLimit); val != "" {
		if err := fscommon.WriteFile(dirPath, "pids.max", val); err != nil {
			return err
		}
	}

	return nil
}

func statPidsWithoutController(dirPath string, stats *cgroups.Stats) error {
	// if the controller is not enabled, let's read PIDS from cgroups.procs
	// (or threads if cgroup.threads is enabled)
	contents, err := fscommon.ReadFile(dirPath, "cgroup.procs")
	if errors.Is(err, unix.ENOTSUP) {
		contents, err = fscommon.ReadFile(dirPath, "cgroup.threads")
	}
	if err != nil {
		return err
	}
	pids := strings.Count(contents, "\n")
	stats.PidsStats.Current = uint64(pids)
	stats.PidsStats.Limit = 0
	return nil
}

func statPids(dirPath string, stats *cgroups.Stats) error {
	current, err := fscommon.GetCgroupParamUint(dirPath, "pids.current")
	if err != nil {
		return errors.Wrap(err, "failed to parse pids.current")
	}

	maxString, err := fscommon.GetCgroupParamString(dirPath, "pids.max")
	if err != nil {
		return errors.Wrap(err, "failed to parse pids.max")
	}

	// Default if pids.max == "max" is 0 -- which represents "no limit".
	var max uint64
	if maxString != "max" {
		max, err = fscommon.ParseUint(maxString, 10, 64)
		if err != nil {
			return errors.Wrapf(err, "failed to parse pids.max - unable to parse %q as a uint from Cgroup file %q",
				maxString, filepath.Join(dirPath, "pids.max"))
		}
	}

	stats.PidsStats.Current = current
	stats.PidsStats.Limit = max
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/rneugeba/virtsock/go/vsock"
	"log"
	"net"
	"os"
	"proxy/libproxy"
	"strings"
)

func onePort() {
	host, port, container := parseHostContainerAddrs()

	vsockP, err := libproxy.NewVsockProxy(&vsock.VsockAddr{Port: uint(port)}, container)
	if err != nil {
		sendError(err)
	}
	ipP, err := libproxy.NewIPProxy(host, container)
	if err != nil {
		sendError(err)
	}

	ctl, err := exposePort(host, port)
	if err != nil {
		sendError(err)
	}

	go handleStopSignals(ipP)
	// TODO: avoid this line if we are running in a TTY
	sendOK()
	go ipP.Run()
	vsockP.Run()
	ctl.Close() // ensure ctl remains alive and un-GCed until here
	os.Exit(0)
}

func exposePort(host net.Addr, port int) (*os.File, error) {
	name := host.Network() + ":" + host.String()
	log.Printf("exposePort %s\n", name)
	err := os.Mkdir("/port/"+name, 0)
	if err != nil {
		log.Printf("Failed to mkdir /port/%s: %#v\n", name, err)
		return nil, err
	}
	ctl, err := os.OpenFile("/port/"+name+"/ctl", os.O_RDWR, 0)
	if err != nil {
		log.Printf("Failed to open /port/%s/ctl: %#v\n", name, err)
		return nil, err
	}
	_, err = ctl.WriteString(fmt.Sprintf("%s:%08x", name, port))
	if err != nil {
		log.Printf("Failed to open /port/%s/ctl: %#v\n", name, err)
		return nil, err
	}
	_, err = ctl.Seek(0, 0)
	if err != nil {
		log.Printf("Failed to seek on /port/%s/ctl: %#v\n", name, err)
		return nil, err
	}
	results := make([]byte, 100)
	count, err := ctl.Read(results)
	if err != nil {
		log.Printf("Failed to read from /port/%s/ctl: %#v\n", name, err)
		return nil, err
	}
	// We deliberately keep the control file open since 9P clunk
	// will trigger a shutdown on the host side.

	response := string(results[0:count])
	if strings.HasPrefix(response, "ERROR ") {
		os.Remove("/port/" + name + "/ctl")
		response = strings.Trim(response[6:], " \t\r\n")
		return nil, errors.New(response)
	}
	// Hold on to a reference to prevent premature GC and close
	return ctl, nil
}
<commit_msg>proxy: update the 9P protocol<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/rneugeba/virtsock/go/vsock"
	"log"
	"net"
	"os"
	"proxy/libproxy"
	"strings"
)

func onePort() {
	host, port, container := parseHostContainerAddrs()

	vsockP, err := libproxy.NewVsockProxy(&vsock.VsockAddr{Port: uint(port)}, container)
	if err != nil {
		sendError(err)
	}
	ipP, err := libproxy.NewIPProxy(host, container)
	if err != nil {
		sendError(err)
	}

	ctl, err := exposePort(host, container)
	if err != nil {
		sendError(err)
	}

	go handleStopSignals(ipP)
	// TODO: avoid this line if we are running in a TTY
	sendOK()
	go ipP.Run()
	vsockP.Run()
	ctl.Close() // ensure ctl remains alive and un-GCed until here
	os.Exit(0)
}

func exposePort(host net.Addr, container net.Addr) (*os.File, error) {
	name := host.Network() + ":" + host.String() + ":" + container.Network() + ":" + container.String()
	log.Printf("exposePort %s\n", name)
	err := os.Mkdir("/port/"+name, 0)
	if err != nil {
		log.Printf("Failed to mkdir /port/%s: %#v\n", name, err)
		return nil, err
	}
	ctl, err := os.OpenFile("/port/"+name+"/ctl", os.O_RDWR, 0)
	if err != nil {
		log.Printf("Failed to open /port/%s/ctl: %#v\n", name, err)
		return nil, err
	}
	_, err = ctl.WriteString(fmt.Sprintf("%s", name))
	if err != nil {
		log.Printf("Failed to open /port/%s/ctl: %#v\n", name, err)
		return nil, err
	}
	_, err = ctl.Seek(0, 0)
	if err != nil {
		log.Printf("Failed to seek on /port/%s/ctl: %#v\n", name, err)
		return nil, err
	}
	results := make([]byte, 100)
	count, err := ctl.Read(results)
	if err != nil {
		log.Printf("Failed to read from /port/%s/ctl: %#v\n", name, err)
		return nil, err
	}
	// We deliberately keep the control file open since 9P clunk
	// will trigger a shutdown on the host side.

	response := string(results[0:count])
	if strings.HasPrefix(response, "ERROR ") {
		os.Remove("/port/" + name + "/ctl")
		response = strings.Trim(response[6:], " \t\r\n")
		return nil, errors.New(response)
	}
	// Hold on to a reference to prevent premature GC and close
	return ctl, nil
}
<|endoftext|>


<commit_before>package importgraph_test

import (
	"go/build"
	"runtime"
	"sort"
	"testing"

	"code.google.com/p/go.tools/refactor/importgraph"

	_ "crypto/hmac" // just for test, below
)

const this = "code.google.com/p/go.tools/refactor/importgraph"

func TestBuild(t *testing.T) {
	saved := runtime.GOMAXPROCS(8) // Build is highly parallel
	defer runtime.GOMAXPROCS(saved)

	forward, reverse, errors := importgraph.Build(&build.Default)

	// Test direct edges.
	// We throw in crypto/hmac to prove that external test files
	// (such as this one) are inspected.
	for _, p := range []string{"go/build", "runtime", "testing", "crypto/hmac"} {
		if !forward[this][p] {
			t.Errorf("forward[importgraph][%s] not found", p)
		}
		if !reverse[p][this] {
			t.Errorf("reverse[%s][importgraph] not found", p)
		}
	}

	// Test non-existent direct edges
	for _, p := range []string{"fmt", "errors", "reflect"} {
		if forward[this][p] {
			t.Errorf("unexpected: forward[importgraph][%s] found", p)
		}
		if reverse[p][this] {
			t.Errorf("unexpected: reverse[%s][importgraph] found", p)
		}
	}

	// Test Search is reflexive.
	if !forward.Search(this)[this] {
		t.Errorf("irreflexive: forward.Search(importgraph)[importgraph] not found")
	}
	if !reverse.Search(this)[this] {
		t.Errorf("irrefexive: reverse.Search(importgraph)[importgraph] not found")
	}

	// Test Search is transitive.  (There is no direct edge to these packages.)
	for _, p := range []string{"errors", "reflect", "unsafe"} {
		if !forward.Search(this)[p] {
			t.Errorf("intransitive: forward.Search(importgraph)[%s] not found", p)
		}
		if !reverse.Search(p)[this] {
			t.Errorf("intransitive: reverse.Search(%s)[importgraph] not found", p)
		}
	}

	// debugging
	if false {
		for path, err := range errors {
			t.Logf("%s: %s", path, err)
		}
		printSorted := func(direction string, g importgraph.Graph, start string) {
			t.Log(direction)
			var pkgs []string
			for pkg := range g.Search(start) {
				pkgs = append(pkgs, pkg)
			}
			sort.Strings(pkgs)
			for _, pkg := range pkgs {
				t.Logf("\t%s", pkg)
			}
		}
		printSorted("forward", forward, this)
		printSorted("forward", reverse, this)
	}
}
<commit_msg>refactor/importgraph: add test of cycles<commit_after>package importgraph_test

import (
	"go/build"
	"runtime"
	"sort"
	"testing"

	"code.google.com/p/go.tools/refactor/importgraph"

	_ "crypto/hmac" // just for test, below
)

const this = "code.google.com/p/go.tools/refactor/importgraph"

func TestBuild(t *testing.T) {
	saved := runtime.GOMAXPROCS(8) // Build is highly parallel
	defer runtime.GOMAXPROCS(saved)

	forward, reverse, errors := importgraph.Build(&build.Default)

	// Test direct edges.
	// We throw in crypto/hmac to prove that external test files
	// (such as this one) are inspected.
	for _, p := range []string{"go/build", "runtime", "testing", "crypto/hmac"} {
		if !forward[this][p] {
			t.Errorf("forward[importgraph][%s] not found", p)
		}
		if !reverse[p][this] {
			t.Errorf("reverse[%s][importgraph] not found", p)
		}
	}

	// Test non-existent direct edges
	for _, p := range []string{"fmt", "errors", "reflect"} {
		if forward[this][p] {
			t.Errorf("unexpected: forward[importgraph][%s] found", p)
		}
		if reverse[p][this] {
			t.Errorf("unexpected: reverse[%s][importgraph] found", p)
		}
	}

	// Test Search is reflexive.
	if !forward.Search(this)[this] {
		t.Errorf("irreflexive: forward.Search(importgraph)[importgraph] not found")
	}
	if !reverse.Search(this)[this] {
		t.Errorf("irrefexive: reverse.Search(importgraph)[importgraph] not found")
	}

	// Test Search is transitive.  (There is no direct edge to these packages.)
	for _, p := range []string{"errors", "reflect", "unsafe"} {
		if !forward.Search(this)[p] {
			t.Errorf("intransitive: forward.Search(importgraph)[%s] not found", p)
		}
		if !reverse.Search(p)[this] {
			t.Errorf("intransitive: reverse.Search(%s)[importgraph] not found", p)
		}
	}

	// Test strongly-connected components.  Because A's external
	// test package can depend on B, and vice versa, most of the
	// standard libraries are mutually dependent when their external
	// tests are considered.
	//
	// For any nodes x, y in the same SCC, y appears in the results
	// of both forward and reverse searches starting from x
	if !forward.Search("fmt")["io"] ||
		!forward.Search("io")["fmt"] ||
		!reverse.Search("fmt")["io"] ||
		!reverse.Search("io")["fmt"] {
		t.Errorf("fmt and io are not mutually reachable despite being in the same SCC")
	}

	// debugging
	if false {
		for path, err := range errors {
			t.Logf("%s: %s", path, err)
		}
		printSorted := func(direction string, g importgraph.Graph, start string) {
			t.Log(direction)
			var pkgs []string
			for pkg := range g.Search(start) {
				pkgs = append(pkgs, pkg)
			}
			sort.Strings(pkgs)
			for _, pkg := range pkgs {
				t.Logf("\t%s", pkg)
			}
		}
		printSorted("forward", forward, this)
		printSorted("reverse", reverse, this)
	}
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"runtime"

	"github.com/cactus/cobra"
)

var Version = "no-version"

func init() {
	RootCmd.AddCommand(&cobra.Command{
		Use:   "version",
		Short: "Print the version",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Printf("rollcage %s (%s,%s-%s)\n", Version,
				runtime.Version(), runtime.Compiler, runtime.GOARCH)
		},
	})
}
<commit_msg>add license output for version<commit_after>package commands

import (
	"fmt"
	"runtime"
	"strings"

	"github.com/cactus/cobra"
)

var Version = "no-version"
var LicenseText = `
This software is available under the MIT License.
    https://github.com/cactus/rollcage

Portions of this software utilize third party libraries:
*   https://github.com/cactus/cobra
    Forked from: https://github.com/spf13/cobra (Apache 2.0 License)
*   https://github.com/cactus/gologit (MIT license)
*   https://github.com/spf13/pflag (BSD license)
*   https://github.com/go-gcfg/gcfg/tree/v1 (BSD license)
`
var showLicense bool

func init() {
	cmd := &cobra.Command{
		Use:   "version",
		Short: "Print the version",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Printf("rollcage %s (%s,%s-%s)\n\n", Version,
				runtime.Version(), runtime.Compiler, runtime.GOARCH)
			if showLicense {
				fmt.Printf("%s\n", strings.TrimSpace(LicenseText))
			}
		},
	}
	cmd.Flags().BoolVarP(
		&showLicense, "license", "l", false,
		"output information about licenses and dependencies")
	RootCmd.AddCommand(cmd)
}
<|endoftext|>