<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/brotherlogic/goserver"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	pbs "github.com/brotherlogic/goserver/proto"
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	runner *Runner
	disk   diskChecker
}

func getHash(file string) (string, error) {
	f, err := os.Open(file)
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return string(h.Sum(nil)), nil
}

func getIP(name string, server string) (string, int) {
	conn, _ := grpc.Dial("192.168.86.34:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	entry := pbd.RegistryEntry{Name: name, Identifier: server}
	r, err := registry.Discover(context.Background(), &entry)

	if err != nil {
		log.Printf("Lookup failed for %v,%v -> %v", name, server, err)
		return "", -1
	}

	return r.Ip, int(r.Port)
}

// updateState of the runner command
func updateState(com *runnerCommand) {
	elems := strings.Split(com.details.Spec.Name, "/")
	dServer, dPort := getIP(elems[len(elems)-1], com.details.Spec.Server)

	if dPort > 0 {
		dConn, err := grpc.Dial(dServer+":"+strconv.Itoa(dPort), grpc.WithInsecure())
		if err != nil {
			panic(err)
		}
		defer dConn.Close()

		c := pbs.NewGoserverServiceClient(dConn)
		_, err = c.IsAlive(context.Background(), &pbs.Alive{})
		com.details.Running = (err == nil)
	} else {
		//Mark as false if we can't locate the job
		com.details.Running = false
	}
}

// BuildJob builds out a job
func (s *Server) BuildJob(ctx context.Context, in *pb.JobSpec) (*pb.Empty, error) {
	s.runner.Checkout(in.Name)
	return &pb.Empty{}, nil
}

// List lists all running jobs
func (s *Server) List(ctx context.Context, in *pb.Empty) (*pb.JobList, error) {
	details := &pb.JobList{}
	for _, job := range s.runner.backgroundTasks {
		updateState(job)
		details.Details = append(details.Details, job.details)
	}

	return details, nil
}

// Run runs a background task
func (s *Server) Run(ctx context.Context, in *pb.JobSpec) (*pb.Empty, error) {
	s.runner.Run(in)
	return &pb.Empty{}, nil
}

//Update restarts a job with new settings
func (s *Server) Update(ctx context.Context, in *pb.JobSpec) (*pb.Empty, error) {
	s.runner.Update(in)
	return &pb.Empty{}, nil
}

// Kill a background task
func (s *Server) Kill(ctx context.Context, in *pb.JobSpec) (*pb.Empty, error) {
	s.runner.kill(in)
	return &pb.Empty{}, nil
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	pb.RegisterGoBuildSlaveServer(server, &s)
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

//Init builds the default runner framework
func Init() *Runner {
	r := &Runner{gopath: "goautobuild"}
	r.runner = runCommand
	go r.run()
	return r
}

func runCommand(c *runnerCommand) {
	log.Printf("RUNNING COMMAND: %v", c)
	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	if len(home) == 0 {

	}
	gpath := home + "/gobuild"
	c.command.Path = strings.Replace(c.command.Path, "$GOPATH", gpath, -1)
	for i := range c.command.Args {
		c.command.Args[i] = strings.Replace(c.command.Args[i], "$GOPATH", gpath, -1)
	}

	path := fmt.Sprintf("GOPATH=" + home + "/gobuild")
	found := false
	log.Printf("HERE = %v", c.command.Env)
	envl := os.Environ()
	for i, blah := range envl {
		if strings.HasPrefix(blah, "GOPATH") {
			envl[i] = path
			found = true
		}
	}
	if !found {
		envl = append(envl, path)
	}
	log.Printf("ENV = %v", envl)
	c.command.Env = envl

	out, err := c.command.StdoutPipe()
	out2, err2 := c.command.StderrPipe()
	if err != nil {
		log.Printf("Blah: %v", err)
	}

	if err2 != nil {
		log.Printf("Blah2: %v", err)
	}

	log.Printf("%v, %v and %v", c.command.Path, c.command.Args, c.command.Env)
	c.command.Start()

	if !c.background {
		buf := new(bytes.Buffer)
		buf.ReadFrom(out)
		str := buf.String()

		buf2 := new(bytes.Buffer)
		buf2.ReadFrom(out2)
		str2 := buf2.String()
		log.Printf("%v and %v", str, str2)

		c.command.Wait()
		c.output = str
		c.complete = true
	}
	log.Printf("DONE")
}

func (diskChecker prodDiskChecker) diskUsage(path string) int64 {
	return diskUsage(path)
}

func (s *Server) rebuildLoop() {
	for true {
		time.Sleep(time.Minute)

		var rebuildList []*pb.JobSpec
		var hashList []string
		for _, job := range s.runner.backgroundTasks {
			log.Printf("Job (started %v, now %v) %v", job.started, time.Now(), job)
			if time.Since(job.started) > time.Hour {
				log.Printf("Added to rebuild list (%v)", job)
				rebuildList = append(rebuildList, job.details.Spec)
				hashList = append(hashList, job.hash)
			}
		}

		log.Printf("Rebuilding %v", rebuildList)
		for i := range rebuildList {
			s.runner.Rebuild(rebuildList[i], hashList[i])
		}
	}
}

func main() {
	s := Server{&goserver.GoServer{}, Init(), prodDiskChecker{}}
	s.Register = s
	s.PrepServer()
	s.RegisterServingTask(s.rebuildLoop)
	s.RegisterServer("gobuildslave", false)
	s.Serve()
}
<commit_msg>Fixed discover<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/brotherlogic/goserver"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	pbs "github.com/brotherlogic/goserver/proto"
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	runner *Runner
	disk   diskChecker
}

func getHash(file string) (string, error) {
	f, err := os.Open(file)
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return string(h.Sum(nil)), nil
}

func getIP(name string, server string) (string, int) {
	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	entry := pbd.RegistryEntry{Name: name, Identifier: server}
	r, err := registry.Discover(context.Background(), &entry)

	if err != nil {
		log.Printf("Lookup failed for %v,%v -> %v", name, server, err)
		return "", -1
	}

	return r.Ip, int(r.Port)
}

// updateState of the runner command
func updateState(com *runnerCommand) {
	elems := strings.Split(com.details.Spec.Name, "/")
	dServer, dPort := getIP(elems[len(elems)-1], com.details.Spec.Server)

	if dPort > 0 {
		dConn, err := grpc.Dial(dServer+":"+strconv.Itoa(dPort), grpc.WithInsecure())
		if err != nil {
			panic(err)
		}
		defer dConn.Close()

		c := pbs.NewGoserverServiceClient(dConn)
		_, err = c.IsAlive(context.Background(), &pbs.Alive{})
		com.details.Running = (err == nil)
	} else {
		//Mark as false if we can't locate the job
		com.details.Running = false
	}
}

// BuildJob builds out a job
func (s *Server) BuildJob(ctx context.Context, in *pb.JobSpec) (*pb.Empty, error) {
	s.runner.Checkout(in.Name)
	return &pb.Empty{}, nil
}

// List lists all running jobs
func (s *Server) List(ctx context.Context, in *pb.Empty) (*pb.JobList, error) {
	details := &pb.JobList{}
	for _, job := range s.runner.backgroundTasks {
		updateState(job)
		details.Details = append(details.Details, job.details)
	}

	return details, nil
}

// Run runs a background task
func (s *Server) Run(ctx context.Context, in *pb.JobSpec) (*pb.Empty, error) {
	s.runner.Run(in)
	return &pb.Empty{}, nil
}

//Update restarts a job with new settings
func (s *Server) Update(ctx context.Context, in *pb.JobSpec) (*pb.Empty, error) {
	s.runner.Update(in)
	return &pb.Empty{}, nil
}

// Kill a background task
func (s *Server) Kill(ctx context.Context, in *pb.JobSpec) (*pb.Empty, error) {
	s.runner.kill(in)
	return &pb.Empty{}, nil
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	pb.RegisterGoBuildSlaveServer(server, &s)
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

//Init builds the default runner framework
func Init() *Runner {
	r := &Runner{gopath: "goautobuild"}
	r.runner = runCommand
	go r.run()
	return r
}

func runCommand(c *runnerCommand) {
	log.Printf("RUNNING COMMAND: %v", c)
	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	if len(home) == 0 {

	}
	gpath := home + "/gobuild"
	c.command.Path = strings.Replace(c.command.Path, "$GOPATH", gpath, -1)
	for i := range c.command.Args {
		c.command.Args[i] = strings.Replace(c.command.Args[i], "$GOPATH", gpath, -1)
	}

	path := fmt.Sprintf("GOPATH=" + home + "/gobuild")
	found := false
	log.Printf("HERE = %v", c.command.Env)
	envl := os.Environ()
	for i, blah := range envl {
		if strings.HasPrefix(blah, "GOPATH") {
			envl[i] = path
			found = true
		}
	}
	if !found {
		envl = append(envl, path)
	}
	log.Printf("ENV = %v", envl)
	c.command.Env = envl

	out, err := c.command.StdoutPipe()
	out2, err2 := c.command.StderrPipe()
	if err != nil {
		log.Printf("Blah: %v", err)
	}

	if err2 != nil {
		log.Printf("Blah2: %v", err)
	}

	log.Printf("%v, %v and %v", c.command.Path, c.command.Args, c.command.Env)
	c.command.Start()

	if !c.background {
		buf := new(bytes.Buffer)
		buf.ReadFrom(out)
		str := buf.String()

		buf2 := new(bytes.Buffer)
		buf2.ReadFrom(out2)
		str2 := buf2.String()
		log.Printf("%v and %v", str, str2)

		c.command.Wait()
		c.output = str
		c.complete = true
	}
	log.Printf("DONE")
}

func (diskChecker prodDiskChecker) diskUsage(path string) int64 {
	return diskUsage(path)
}

func (s *Server) rebuildLoop() {
	for true {
		time.Sleep(time.Minute)

		var rebuildList []*pb.JobSpec
		var hashList []string
		for _, job := range s.runner.backgroundTasks {
			log.Printf("Job (started %v, now %v) %v", job.started, time.Now(), job)
			if time.Since(job.started) > time.Hour {
				log.Printf("Added to rebuild list (%v)", job)
				rebuildList = append(rebuildList, job.details.Spec)
				hashList = append(hashList, job.hash)
			}
		}

		log.Printf("Rebuilding %v", rebuildList)
		for i := range rebuildList {
			s.runner.Rebuild(rebuildList[i], hashList[i])
		}
	}
}

func main() {
	s := Server{&goserver.GoServer{}, Init(), prodDiskChecker{}}
	s.Register = s
	s.PrepServer()
	s.RegisterServingTask(s.rebuildLoop)
	s.RegisterServer("gobuildslave", false)
	s.Serve()
}
<|endoftext|>


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

package jsonrpc2

import (
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"time"

	"golang.org/x/tools/internal/event"
	errors "golang.org/x/xerrors"
)

// NOTE: This file provides an experimental API for serving multiple remote
// jsonrpc2 clients over the network. For now, it is intentionally similar to
// net/http, but that may change in the future as we figure out the correct
// semantics.

// A StreamServer is used to serve incoming jsonrpc2 clients communicating over
// a newly created connection.
type StreamServer interface {
	ServeStream(context.Context, Conn) error
}

// The ServerFunc type is an adapter that implements the StreamServer interface
// using an ordinary function.
type ServerFunc func(context.Context, Conn) error

// ServeStream calls f(ctx, s).
func (f ServerFunc) ServeStream(ctx context.Context, c Conn) error {
	return f(ctx, c)
}

// HandlerServer returns a StreamServer that handles incoming streams using the
// provided handler.
func HandlerServer(h Handler) StreamServer {
	return ServerFunc(func(ctx context.Context, conn Conn) error {
		conn.Go(ctx, h)
		<-conn.Done()
		return conn.Err()
	})
}

// ListenAndServe starts an jsonrpc2 server on the given address.  If
// idleTimeout is non-zero, ListenAndServe exits after there are no clients for
// this duration, otherwise it exits only on error.
func ListenAndServe(ctx context.Context, network, addr string, server StreamServer, idleTimeout time.Duration) error {
	ln, err := net.Listen(network, addr)
	if err != nil {
		return err
	}
	defer ln.Close()
	if network == "unix" {
		defer os.Remove(addr)
	}
	return Serve(ctx, ln, server, idleTimeout)
}

// Serve accepts incoming connections from the network, and handles them using
// the provided server. If idleTimeout is non-zero, ListenAndServe exits after
// there are no clients for this duration, otherwise it exits only on error.
func Serve(ctx context.Context, ln net.Listener, server StreamServer, idleTimeout time.Duration) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	// Max duration: ~290 years; surely that's long enough.
	const forever = 1<<63 - 1
	if idleTimeout <= 0 {
		idleTimeout = forever
	}
	connTimer := time.NewTimer(idleTimeout)

	newConns := make(chan net.Conn)
	doneListening := make(chan error)
	closedConns := make(chan error)

	go func() {
		for {
			nc, err := ln.Accept()
			if err != nil {
				select {
				case doneListening <- fmt.Errorf("Accept(): %w", err):
				case <-ctx.Done():
				}
				return
			}
			newConns <- nc
		}
	}()

	activeConns := 0
	for {
		select {
		case netConn := <-newConns:
			activeConns++
			connTimer.Stop()
			stream := NewHeaderStream(netConn)
			go func() {
				conn := NewConn(stream)
				closedConns <- server.ServeStream(ctx, conn)
				stream.Close()
			}()
		case err := <-doneListening:
			return err
		case err := <-closedConns:
			if !isClosingError(err) {
				event.Error(ctx, "closed a connection", err)
			}
			activeConns--
			if activeConns == 0 {
				connTimer.Reset(idleTimeout)
			}
		case <-connTimer.C:
			return ErrIdleTimeout
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// isClosingError reports if the error occurs normally during the process of
// closing a network connection. It uses imperfect heuristics that err on the
// side of false negatives, and should not be used for anything critical.
func isClosingError(err error) bool {
	if errors.Is(err, io.EOF) {
		return true
	}
	// Per https://github.com/golang/go/issues/4373, this error string should not
	// change. This is not ideal, but since the worst that could happen here is
	// some superfluous logging, it is acceptable.
	if err.Error() == "use of closed network connection" {
		return true
	}
	return false
}
<commit_msg>internal/jsonrpc2: make Serve wait for all connections to close<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package jsonrpc2

import (
	"context"
	"io"
	"net"
	"os"
	"time"

	"golang.org/x/tools/internal/event"
	errors "golang.org/x/xerrors"
)

// NOTE: This file provides an experimental API for serving multiple remote
// jsonrpc2 clients over the network. For now, it is intentionally similar to
// net/http, but that may change in the future as we figure out the correct
// semantics.

// A StreamServer is used to serve incoming jsonrpc2 clients communicating over
// a newly created connection.
type StreamServer interface {
	ServeStream(context.Context, Conn) error
}

// The ServerFunc type is an adapter that implements the StreamServer interface
// using an ordinary function.
type ServerFunc func(context.Context, Conn) error

// ServeStream calls f(ctx, s).
func (f ServerFunc) ServeStream(ctx context.Context, c Conn) error {
	return f(ctx, c)
}

// HandlerServer returns a StreamServer that handles incoming streams using the
// provided handler.
func HandlerServer(h Handler) StreamServer {
	return ServerFunc(func(ctx context.Context, conn Conn) error {
		conn.Go(ctx, h)
		<-conn.Done()
		return conn.Err()
	})
}

// ListenAndServe starts an jsonrpc2 server on the given address.  If
// idleTimeout is non-zero, ListenAndServe exits after there are no clients for
// this duration, otherwise it exits only on error.
func ListenAndServe(ctx context.Context, network, addr string, server StreamServer, idleTimeout time.Duration) error {
	ln, err := net.Listen(network, addr)
	if err != nil {
		return err
	}
	defer ln.Close()
	if network == "unix" {
		defer os.Remove(addr)
	}
	return Serve(ctx, ln, server, idleTimeout)
}

// Serve accepts incoming connections from the network, and handles them using
// the provided server. If idleTimeout is non-zero, ListenAndServe exits after
// there are no clients for this duration, otherwise it exits only on error.
func Serve(ctx context.Context, ln net.Listener, server StreamServer, idleTimeout time.Duration) error {
	newConns := make(chan net.Conn)
	closedConns := make(chan error)
	activeConns := 0
	var acceptErr error
	go func() {
		defer close(newConns)
		for {
			var nc net.Conn
			nc, acceptErr = ln.Accept()
			if acceptErr != nil {
				return
			}
			newConns <- nc
		}
	}()

	ctx, cancel := context.WithCancel(ctx)
	defer func() {
		// Signal the Accept goroutine to stop immediately
		// and terminate all newly-accepted connections until it returns.
		ln.Close()
		for nc := range newConns {
			nc.Close()
		}
		// Cancel pending ServeStream callbacks and wait for them to finish.
		cancel()
		for activeConns > 0 {
			err := <-closedConns
			if !isClosingError(err) {
				event.Error(ctx, "closed a connection", err)
			}
			activeConns--
		}
	}()

	// Max duration: ~290 years; surely that's long enough.
	const forever = 1<<63 - 1
	if idleTimeout <= 0 {
		idleTimeout = forever
	}
	connTimer := time.NewTimer(idleTimeout)
	defer connTimer.Stop()

	for {
		select {
		case netConn, ok := <-newConns:
			if !ok {
				return acceptErr
			}
			if activeConns == 0 && !connTimer.Stop() {
				// connTimer.C may receive a value even after Stop returns.
				// (See https://golang.org/issue/37196.)
				<-connTimer.C
			}
			activeConns++
			stream := NewHeaderStream(netConn)
			go func() {
				conn := NewConn(stream)
				err := server.ServeStream(ctx, conn)
				stream.Close()
				closedConns <- err
			}()

		case err := <-closedConns:
			if !isClosingError(err) {
				event.Error(ctx, "closed a connection", err)
			}
			activeConns--
			if activeConns == 0 {
				connTimer.Reset(idleTimeout)
			}

		case <-connTimer.C:
			return ErrIdleTimeout

		case <-ctx.Done():
			return nil
		}
	}
}

// isClosingError reports if the error occurs normally during the process of
// closing a network connection. It uses imperfect heuristics that err on the
// side of false negatives, and should not be used for anything critical.
func isClosingError(err error) bool {
	if errors.Is(err, io.EOF) {
		return true
	}
	// Per https://github.com/golang/go/issues/4373, this error string should not
	// change. This is not ideal, but since the worst that could happen here is
	// some superfluous logging, it is acceptable.
	if err.Error() == "use of closed network connection" {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>/*

	patchtree is a simple library that knows how to interpret a folder
	structure of jd diffs and apply them on top of a base.

		test/
		  base.json
		  after_move/
		    modification.patch
		  sanitization/
		    modification.patch
		    hidden/
		      modification.patch
		    nonempty/
		      modification.patch

	Given a path relative ot the current binary, it walks backwards up the
	folder, ensuring that a modification.patch exists in each directory until
	it finds a base.json. Then it applies forward all of the
	modification.patches to give you the final composed json blob result.

	patchtree-helper is a simple cli utility that wraps the functionality in
	this package.

*/
package patchtree

import (
	"encoding/json"
	"errors"
	jd "github.com/jkomoros/jd/lib"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

//BATCH_JSON_NAME is the name of the base json file at the root of the
//directory tree.
const BASE_JSON_NAME = "base.json"

//PATCH_NAME is the name of the patch in each sub-folder that modifies the
//json in the tree above it.
const PATCH_NAME = "modification.patch"

//EXPANDED_JSON_NAME is the name of the file that represents the entire
//expanded json blob at a given part of the tree, created by `expand`.
const EXPANDED_JSON_NAME = "node.expanded.json"

//JSON returns the patched json blob impplied by that directory structure or
//an error if something doesn't work. See the package doc for more.
func JSON(path string) ([]byte, error) {

	if strings.HasSuffix(path, "/") {
		path = strings.TrimSuffix(path, "/")
	}

	result, err := processDirectory(path)

	if err != nil {
		return nil, err
	}

	return []byte(result.Json()), nil

}

//MustJSON is the same as JSON, but if it would have returned an error, panics istead.
func MustJSON(path string) []byte {
	result, err := JSON(path)
	if err != nil {
		panic(err)
	}
	return result
}

type directoryFunc func(string, jd.JsonNode) (int, error)

func startDirectoryAndWalk(rootPath string, subFunc directoryFunc) (int, error) {
	baseJsonPath := filepath.Clean(rootPath + "/" + BASE_JSON_NAME)

	if _, err := os.Stat(baseJsonPath); os.IsNotExist(err) {
		return 0, errors.New("Base json file did not exist: " + err.Error())
	}

	node, err := jd.ReadJsonFile(baseJsonPath)

	if err != nil {
		return 0, errors.New("Couldn't parse base json file: " + err.Error())
	}

	return walkDirectory(rootPath, node, subFunc)
}

func walkDirectory(directory string, expandedNode jd.JsonNode, subFunc directoryFunc) (int, error) {
	files, err := ioutil.ReadDir(directory)

	if err != nil {
		return 0, errors.New("Couldn't read directory: " + err.Error())
	}

	numAffectedFiles := 0

	for _, file := range files {
		if !file.IsDir() {
			continue
		}
		subDirectory := filepath.Clean(directory + "/" + file.Name())

		subAffectedFiles, err := subFunc(subDirectory, expandedNode)
		if err != nil {
			return numAffectedFiles, err
		}
		numAffectedFiles += subAffectedFiles
	}

	return numAffectedFiles, nil
}

//ExpandTree expands all of the nodes in the patchtree, applying the chains of
//modification and created an node.expanded.json in each node. Used in a
//workflow to modify base.json: run this commeand, then modify base.json, then
//run ContractTree.
func ExpandTree(rootPath string) (affectedFiles int, err error) {
	return startDirectoryAndWalk(rootPath, expandTreeProcessDirectory)
}

func expandTreeProcessDirectory(directory string, node jd.JsonNode) (int, error) {

	diffFileName := filepath.Clean(directory + "/" + PATCH_NAME)

	if _, err := os.Stat(diffFileName); os.IsNotExist(err) {
		//TODO: it's weird to print to log when this condition is hit.
		log.Println(diffFileName + " did not exist; skipping that directory and all beneath it.")
		return 0, nil
	}

	diff, err := jd.ReadDiffFile(diffFileName)

	if err != nil {
		return 0, errors.New(diffFileName + " could not be loaded as patch file: " + err.Error())
	}

	expandedNode, err := node.Patch(diff)

	if err != nil {
		return 0, errors.New(diffFileName + " could not be applied: " + err.Error())
	}

	expandedNodeFileName := filepath.Clean(directory + "/" + EXPANDED_JSON_NAME)

	data := expandedNode.Json()

	indentedJson, err := indentJson(data)

	if err != nil {
		return 0, errors.New("Couldn't indent json: " + err.Error())
	}

	if err := ioutil.WriteFile(expandedNodeFileName, []byte(indentedJson), 0644); err != nil {
		return 0, errors.New("Couldn't write " + expandedNodeFileName + ": " + err.Error())
	}

	numAffectedFiles, err := walkDirectory(directory, expandedNode, expandTreeProcessDirectory)

	if err != nil {
		return 0, err
	}

	return numAffectedFiles + 1, nil

}

func indentJson(data string) (indented string, err error) {
	var obj map[string]interface{}

	if err := json.Unmarshal([]byte(data), &obj); err != nil {
		return "", errors.New("Couldn't unpack generated json: " + err.Error())
	}

	result, err := json.MarshalIndent(obj, "", "\t")

	if err != nil {
		return "", errors.New("Couldn't repack generated json: " + err.Error())
	}

	return string(result), nil

}

//ContractTree goes through each node in the parse tree and where it finds a
//node.expanded,json, re-derives and overwrites the "modification.patch". Used
//as part of a workflow to modify base.json: run ExpandTree, modify base.json,
//then ContractTree.
func ContractTree(rootPath string) (numAffectedFiles int, err error) {
	return startDirectoryAndWalk(rootPath, contractTreeProcessDirectory)
}

func contractTreeProcessDirectory(directory string, node jd.JsonNode) (int, error) {

	nodeFileName := filepath.Clean(directory + "/" + EXPANDED_JSON_NAME)

	expandedNode, err := jd.ReadJsonFile(nodeFileName)

	if err != nil {
		return 0, errors.New(nodeFileName + " could not be loaded as json file: " + err.Error())
	}

	patch := node.Diff(expandedNode)

	data := patch.Render()

	diffFileName := filepath.Clean(directory + "/" + PATCH_NAME)

	if err := ioutil.WriteFile(diffFileName, []byte(data), 0644); err != nil {
		return 0, errors.New("Couldn't write " + diffFileName + ": " + err.Error())
	}

	numAffectedFiles, err := walkDirectory(directory, expandedNode, contractTreeProcessDirectory)

	if err != nil {
		return 0, err
	}

	return numAffectedFiles + 1, nil

}

//CleanTree goes through each node, and if modification.json conceptually
//matches node.expanded.json, then removes node.expanded.json.
func CleanTree(rootPath string) (numAffectedFiles int, err error) {
	return startDirectoryAndWalk(rootPath, cleanTreeProcessDirectory)
}

func cleanTreeProcessDirectory(directory string, node jd.JsonNode) (int, error) {

	nodeFileName := filepath.Clean(directory + "/" + EXPANDED_JSON_NAME)

	expandedNode, err := jd.ReadJsonFile(nodeFileName)

	if err != nil {
		return 0, errors.New(nodeFileName + " could not be loaded as json file: " + err.Error())
	}

	patchFileName := filepath.Clean(directory + "/" + PATCH_NAME)

	patch, err := jd.ReadDiffFile(patchFileName)

	if err != nil {
		return 0, errors.New(patchFileName + " could not be loaded as diff file: " + err.Error())
	}

	expandedNodeWithPatch, err := node.Patch(patch)

	if err != nil {
		return 0, errors.New(directory + " patch could not be applied: " + err.Error())
	}

	if !expandedNodeWithPatch.Equals(expandedNode) {
		return 0, errors.New(directory + " patch file did not match expanded json")
	}

	if err := os.Remove(nodeFileName); err != nil {
		return 0, errors.New("Couldn't delete node file: " + err.Error())
	}

	numAffectedFiles, err := walkDirectory(directory, expandedNode, cleanTreeProcessDirectory)

	if err != nil {
		return 0, err
	}

	return numAffectedFiles + 1, nil

}

func processDirectory(path string) (jd.JsonNode, error) {

	//If no more path pieces error
	if path == "" || path == "/" || path == "./" {
		return nil, errors.New("Didn't find a base.json anywhere in the given directory structure")
	}

	//TODO: check if the directory exists...

	baseJsonPath := filepath.Clean(path + "/" + BASE_JSON_NAME)

	if _, err := os.Stat(baseJsonPath); err == nil {
		//Found the directory with base.json!
		node, err := jd.ReadJsonFile(baseJsonPath)
		if err != nil {
			return nil, errors.New(path + " had error reading base.json: " + err.Error())
		}
		return node, nil
	}

	modificationPatchPath := filepath.Clean(path + "/" + PATCH_NAME)

	if _, err := os.Stat(modificationPatchPath); err == nil {

		//Recurse, with the sub-directory.
		baseJson, err := processDirectory(filepath.Dir(path))

		if err != nil {
			return nil, err
		}

		diff, err := jd.ReadDiffFile(modificationPatchPath)

		if err != nil {
			return nil, errors.New("Error reading diff file at " + modificationPatchPath + ": " + err.Error())
		}

		composed, err := baseJson.Patch(diff)

		if err != nil {
			return nil, errors.New(path + " had error diffing " + err.Error())
		}

		return composed, nil
	}

	//Path had neither base.json or modification.patch, which is an error
	return nil, errors.New("In " + path + " didn't have either " + BASE_JSON_NAME + " or " + PATCH_NAME)

}
<commit_msg>Fix the bug where sibling patches don't work, by copying json before calling walkDirectory on each node. Fixes #608.<commit_after>/*

	patchtree is a simple library that knows how to interpret a folder
	structure of jd diffs and apply them on top of a base.

		test/
		  base.json
		  after_move/
		    modification.patch
		  sanitization/
		    modification.patch
		    hidden/
		      modification.patch
		    nonempty/
		      modification.patch

	Given a path relative ot the current binary, it walks backwards up the
	folder, ensuring that a modification.patch exists in each directory until
	it finds a base.json. Then it applies forward all of the
	modification.patches to give you the final composed json blob result.

	patchtree-helper is a simple cli utility that wraps the functionality in
	this package.

*/
package patchtree

import (
	"encoding/json"
	"errors"
	jd "github.com/jkomoros/jd/lib"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

//BATCH_JSON_NAME is the name of the base json file at the root of the
//directory tree.
const BASE_JSON_NAME = "base.json"

//PATCH_NAME is the name of the patch in each sub-folder that modifies the
//json in the tree above it.
const PATCH_NAME = "modification.patch"

//EXPANDED_JSON_NAME is the name of the file that represents the entire
//expanded json blob at a given part of the tree, created by `expand`.
const EXPANDED_JSON_NAME = "node.expanded.json"

//JSON returns the patched json blob impplied by that directory structure or
//an error if something doesn't work. See the package doc for more.
func JSON(path string) ([]byte, error) {

	if strings.HasSuffix(path, "/") {
		path = strings.TrimSuffix(path, "/")
	}

	result, err := processDirectory(path)

	if err != nil {
		return nil, err
	}

	return []byte(result.Json()), nil

}

//MustJSON is the same as JSON, but if it would have returned an error, panics istead.
func MustJSON(path string) []byte {
	result, err := JSON(path)
	if err != nil {
		panic(err)
	}
	return result
}

type directoryFunc func(string, jd.JsonNode) (int, error)

func startDirectoryAndWalk(rootPath string, subFunc directoryFunc) (int, error) {
	baseJsonPath := filepath.Clean(rootPath + "/" + BASE_JSON_NAME)

	if _, err := os.Stat(baseJsonPath); os.IsNotExist(err) {
		return 0, errors.New("Base json file did not exist: " + err.Error())
	}

	node, err := jd.ReadJsonFile(baseJsonPath)

	if err != nil {
		return 0, errors.New("Couldn't parse base json file: " + err.Error())
	}

	return walkDirectory(rootPath, node, subFunc)
}

func walkDirectory(directory string, expandedNode jd.JsonNode, subFunc directoryFunc) (int, error) {
	files, err := ioutil.ReadDir(directory)

	if err != nil {
		return 0, errors.New("Couldn't read directory: " + err.Error())
	}

	numAffectedFiles := 0

	for _, file := range files {
		if !file.IsDir() {
			continue
		}
		subDirectory := filepath.Clean(directory + "/" + file.Name())

		subAffectedFiles, err := subFunc(subDirectory, copyJson(expandedNode))
		if err != nil {
			return numAffectedFiles, err
		}
		numAffectedFiles += subAffectedFiles
	}

	return numAffectedFiles, nil
}

//copyJson returns a copy of the given json node. This is necessary because
//methods like Patch() actually modify the underlying json node (even though
//that's unclear!)
func copyJson(node jd.JsonNode) jd.JsonNode {
	result, _ := jd.ReadJsonString(node.Json())

	return result
}

//ExpandTree expands all of the nodes in the patchtree, applying the chains of
//modification and created an node.expanded.json in each node. Used in a
//workflow to modify base.json: run this commeand, then modify base.json, then
//run ContractTree.
func ExpandTree(rootPath string) (affectedFiles int, err error) {
	return startDirectoryAndWalk(rootPath, expandTreeProcessDirectory)
}

func expandTreeProcessDirectory(directory string, node jd.JsonNode) (int, error) {

	diffFileName := filepath.Clean(directory + "/" + PATCH_NAME)

	if _, err := os.Stat(diffFileName); os.IsNotExist(err) {
		//TODO: it's weird to print to log when this condition is hit.
		log.Println(diffFileName + " did not exist; skipping that directory and all beneath it.")
		return 0, nil
	}

	diff, err := jd.ReadDiffFile(diffFileName)

	if err != nil {
		return 0, errors.New(diffFileName + " could not be loaded as patch file: " + err.Error())
	}

	expandedNode, err := node.Patch(diff)

	if err != nil {
		return 0, errors.New(diffFileName + " could not be applied: " + err.Error())
	}

	expandedNodeFileName := filepath.Clean(directory + "/" + EXPANDED_JSON_NAME)

	data := expandedNode.Json()

	indentedJson, err := indentJson(data)

	if err != nil {
		return 0, errors.New("Couldn't indent json: " + err.Error())
	}

	if err := ioutil.WriteFile(expandedNodeFileName, []byte(indentedJson), 0644); err != nil {
		return 0, errors.New("Couldn't write " + expandedNodeFileName + ": " + err.Error())
	}

	numAffectedFiles, err := walkDirectory(directory, expandedNode, expandTreeProcessDirectory)

	if err != nil {
		return 0, err
	}

	return numAffectedFiles + 1, nil

}

func indentJson(data string) (indented string, err error) {
	var obj map[string]interface{}

	if err := json.Unmarshal([]byte(data), &obj); err != nil {
		return "", errors.New("Couldn't unpack generated json: " + err.Error())
	}

	result, err := json.MarshalIndent(obj, "", "\t")

	if err != nil {
		return "", errors.New("Couldn't repack generated json: " + err.Error())
	}

	return string(result), nil

}

//ContractTree goes through each node in the parse tree and where it finds a
//node.expanded,json, re-derives and overwrites the "modification.patch". Used
//as part of a workflow to modify base.json: run ExpandTree, modify base.json,
//then ContractTree.
func ContractTree(rootPath string) (numAffectedFiles int, err error) {
	return startDirectoryAndWalk(rootPath, contractTreeProcessDirectory)
}

func contractTreeProcessDirectory(directory string, node jd.JsonNode) (int, error) {

	nodeFileName := filepath.Clean(directory + "/" + EXPANDED_JSON_NAME)

	expandedNode, err := jd.ReadJsonFile(nodeFileName)

	if err != nil {
		return 0, errors.New(nodeFileName + " could not be loaded as json file: " + err.Error())
	}

	patch := node.Diff(expandedNode)

	data := patch.Render()

	diffFileName := filepath.Clean(directory + "/" + PATCH_NAME)

	if err := ioutil.WriteFile(diffFileName, []byte(data), 0644); err != nil {
		return 0, errors.New("Couldn't write " + diffFileName + ": " + err.Error())
	}

	numAffectedFiles, err := walkDirectory(directory, expandedNode, contractTreeProcessDirectory)

	if err != nil {
		return 0, err
	}

	return numAffectedFiles + 1, nil

}

//CleanTree goes through each node, and if modification.json conceptually
//matches node.expanded.json, then removes node.expanded.json.
func CleanTree(rootPath string) (numAffectedFiles int, err error) {
	return startDirectoryAndWalk(rootPath, cleanTreeProcessDirectory)
}

func cleanTreeProcessDirectory(directory string, node jd.JsonNode) (int, error) {

	nodeFileName := filepath.Clean(directory + "/" + EXPANDED_JSON_NAME)

	expandedNode, err := jd.ReadJsonFile(nodeFileName)

	if err != nil {
		return 0, errors.New(nodeFileName + " could not be loaded as json file: " + err.Error())
	}

	patchFileName := filepath.Clean(directory + "/" + PATCH_NAME)

	patch, err := jd.ReadDiffFile(patchFileName)

	if err != nil {
		return 0, errors.New(patchFileName + " could not be loaded as diff file: " + err.Error())
	}

	expandedNodeWithPatch, err := node.Patch(patch)

	if err != nil {
		return 0, errors.New(directory + " patch could not be applied: " + err.Error())
	}

	if !expandedNodeWithPatch.Equals(expandedNode) {
		return 0, errors.New(directory + " patch file did not match expanded json")
	}

	if err := os.Remove(nodeFileName); err != nil {
		return 0, errors.New("Couldn't delete node file: " + err.Error())
	}

	numAffectedFiles, err := walkDirectory(directory, expandedNode, cleanTreeProcessDirectory)

	if err != nil {
		return 0, err
	}

	return numAffectedFiles + 1, nil

}

func processDirectory(path string) (jd.JsonNode, error) {

	//If no more path pieces error
	if path == "" || path == "/" || path == "./" {
		return nil, errors.New("Didn't find a base.json anywhere in the given directory structure")
	}

	//TODO: check if the directory exists...

	baseJsonPath := filepath.Clean(path + "/" + BASE_JSON_NAME)

	if _, err := os.Stat(baseJsonPath); err == nil {
		//Found the directory with base.json!
		node, err := jd.ReadJsonFile(baseJsonPath)
		if err != nil {
			return nil, errors.New(path + " had error reading base.json: " + err.Error())
		}
		return node, nil
	}

	modificationPatchPath := filepath.Clean(path + "/" + PATCH_NAME)

	if _, err := os.Stat(modificationPatchPath); err == nil {

		//Recurse, with the sub-directory.
		baseJson, err := processDirectory(filepath.Dir(path))

		if err != nil {
			return nil, err
		}

		diff, err := jd.ReadDiffFile(modificationPatchPath)

		if err != nil {
			return nil, errors.New("Error reading diff file at " + modificationPatchPath + ": " + err.Error())
		}

		composed, err := baseJson.Patch(diff)

		if err != nil {
			return nil, errors.New(path + " had error diffing " + err.Error())
		}

		return composed, nil
	}

	//Path had neither base.json or modification.patch, which is an error
	return nil, errors.New("In " + path + " didn't have either " + BASE_JSON_NAME + " or " + PATCH_NAME)

}
<|endoftext|>


<commit_before>package version

import (
	"fmt"
	"runtime"
)

const Binary = "0.3.7"

func String(app string) string {
	return fmt.Sprintf("%s v%s (built w/%s)", app, Binary, runtime.Version())
}
<commit_msg>bump v0.3.8-rc1<commit_after>package version

import (
	"fmt"
	"runtime"
)

const Binary = "0.3.8-rc1"

func String(app string) string {
	return fmt.Sprintf("%s v%s (built w/%s)", app, Binary, runtime.Version())
}
<|endoftext|>


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

//+build linux

package main

// #cgo LDFLAGS: -ldl
// #include <dlfcn.h>
// #include <sys/types.h>
//
// int
// my_sd_pid_get_owner_uid(void *f, pid_t pid, uid_t *uid)
// {
//   int (*sd_pid_get_owner_uid)(pid_t, uid_t *);
//
//   sd_pid_get_owner_uid = (int (*)(pid_t, uid_t *))f;
//   return sd_pid_get_owner_uid(pid, uid);
// }
//
import "C"

// this implements /init of stage1/nspawn+systemd

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"syscall"
	"strings"

	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/appc/spec/schema/types"

	"github.com/coreos/rkt/common"
	"github.com/coreos/rkt/networking"
	"github.com/coreos/rkt/pkg/sys"
)

const (
	// Path to systemd-nspawn binary within the stage1 rootfs
	nspawnBin = "/usr/bin/systemd-nspawn"
	// Path to lkvm binary within the stage1 rootfs
	lkvmBin = "/usr/bin/lkvm"
	bzImg = "/usr/lib/kernel/vmlinux.container"
	// Path to the interpreter within the stage1 rootfs
	interpBin = "/usr/lib64/ld-linux-x86-64.so.2"
	// Path to the localtime file/symlink in host
	localtimePath = "/etc/localtime"
)

// mirrorLocalZoneInfo tries to reproduce the /etc/localtime target in stage1/ to satisfy systemd-nspawn
func mirrorLocalZoneInfo(root string) {
	zif, err := os.Readlink(localtimePath)
	if err != nil {
		return
	}

	// On some systems /etc/localtime is a relative symlink, make it absolute
	if !filepath.IsAbs(zif) {
		zif = filepath.Join(filepath.Dir(localtimePath), zif)
		zif = filepath.Clean(zif)
	}

	src, err := os.Open(zif)
	if err != nil {
		return
	}
	defer src.Close()

	destp := filepath.Join(common.Stage1RootfsPath(root), zif)

	if err = os.MkdirAll(filepath.Dir(destp), 0755); err != nil {
		return
	}

	dest, err := os.OpenFile(destp, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer dest.Close()

	_, _ = io.Copy(dest, src)
}

var (
	debug       bool
	privNet     bool
	interactive bool
	virtualisation string
)

func init() {
	flag.BoolVar(&debug, "debug", false, "Run in debug mode")
	flag.BoolVar(&privNet, "private-net", false, "Setup private network")
	flag.BoolVar(&interactive, "interactive", false, "The pod is interactive")
	flag.StringVar(&virtualisation, "containment-type", "kvm", "Containment type to use: nspawn or kvm (default)")

	if os.Getenv("RKT_CONTAINMENT_TYPE") != "" {
		virtualisation = os.Getenv("RKT_CONTAINMENT_TYPE")
	}

	// this ensures that main runs only on main thread (thread group leader).
	// since namespace ops (unshare, setns) are done for a single thread, we
	// must ensure that the goroutine does not jump from OS thread to thread
	runtime.LockOSThread()
}

// getArgsEnvNspawn returns the nspawn args and env according to the usr used
func getArgsEnvNspawn(p *Pod) ([]string, []string, error) {
	args := []string{}
	env := os.Environ()

	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), interpBin))
	args = append(args, "--library-path")
	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), "usr/lib64"))

	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), nspawnBin))
	args = append(args, "--boot") // Launch systemd in the pod
	out, err := os.Getwd()
	if err != nil {
		return nil, nil, err
	}
	lfd, err := common.GetRktLockFD()
	if err != nil {
		return nil, nil, err
	}
	args = append(args, fmt.Sprintf("--pid-file=%v", filepath.Join(out, "pid")))
	args = append(args, fmt.Sprintf("--keep-fd=%v", lfd))
	args = append(args, fmt.Sprintf("--register=true"))

	if !debug {
		args = append(args, "--quiet") // silence most nspawn output (log_warning is currently not covered by this)
	}

	keepUnit, err := runningFromUnitFile()
	if err != nil {
		return nil, nil, fmt.Errorf("error determining if we're running from a unit file: %v", err)
	}

	if keepUnit {
		args = append(args, "--keep-unit")
	}

	nsargs, err := p.PodToNspawnArgs()
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to generate nspawn args: %v", err)
	}
	args = append(args, nsargs...)

	args = append(args, "--")
	args = append(args, "--default-standard-output=tty")

	if !debug {
		args = append(args, "--log-target=null")
		args = append(args, "--show-status=0")
	}

	return args, env, nil
}

func getArgsEnvKvm(p *Pod) ([]string, []string, error) {
	args := []string{}
	kargs := []string{}
	env := os.Environ()

	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), interpBin))
	args = append(args, "--library-path")
	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), "usr/lib64"))

	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), lkvmBin))
	args = append(args, "run")

	args = append(args, "-m 1024")
	args = append(args, "-c 6")

	args = append(args, fmt.Sprintf("--kernel=%v", filepath.Join(common.Stage1RootfsPath(p.Root), bzImg)))
	args = append(args, "--console=virtio")
	kargs = append(kargs, "console=hvc0")

	kargs = append(kargs, "init=/usr/lib/systemd/systemd")

	nsargs, err := p.PodToKvmArgs()
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to generate kvm args: %v", err)
	}
	args = append(args, nsargs...)

	// Arguments to systemd
	kargs = append(kargs, "systemd.default_standard_output=tty")
	if !debug {
		kargs = append(kargs, "systemd.log_target=null")
		kargs = append(kargs, "systemd.show-status=0")
		kargs = append(kargs, "quiet") // silence most nspawn output (log_warning is currently not covered by this)
	}

	args = append(args, "--param")
	args = append(args, strings.Join(kargs, " "))

	return args, env, nil
}

func getArgsEnv(p *Pod) ([]string, []string, error) {
	switch virtualisation {
	case "nspawn":
		return getArgsEnvNspawn(p)
	case "kvm":
		return getArgsEnvKvm(p)
	default:
		return nil, nil, fmt.Errorf("unrecognized containment type: %v", virtualisation)
	}
}

func withClearedCloExec(lfd int, f func() error) error {
	err := sys.CloseOnExec(lfd, false)
	if err != nil {
		return err
	}
	defer sys.CloseOnExec(lfd, true)

	return f()
}

func forwardedPorts(pod *Pod) ([]networking.ForwardedPort, error) {
	fps := []networking.ForwardedPort{}

	for _, ep := range pod.Manifest.Ports {
		n := ""
		fp := networking.ForwardedPort{}

		for _, a := range pod.Manifest.Apps {
			for _, p := range a.App.Ports {
				if p.Name == ep.Name {
					if n == "" {
						fp.Protocol = p.Protocol
						fp.HostPort = ep.HostPort
						fp.PodPort = p.Port
						n = a.Name.String()
					} else {
						return nil, fmt.Errorf("Ambiguous exposed port in PodManifest: %q and %q both define port %q", n, a.Name, p.Name)
					}
				}
			}
		}

		if n == "" {
			return nil, fmt.Errorf("Port name %q is not defined by any apps", ep.Name)
		}

		fps = append(fps, fp)
	}

	// TODO(eyakubovich): validate that there're no conflicts

	return fps, nil
}

func stage1() int {
	uuid, err := types.NewUUID(flag.Arg(0))
	if err != nil {
		fmt.Fprintln(os.Stderr, "UUID is missing or malformed")
		return 1
	}

	root := "."
	p, err := LoadPod(root, uuid)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load pod: %v\n", err)
		return 1
	}

	// set close-on-exec flag on RKT_LOCK_FD so it gets correctly closed when invoking
	// network plugins
	lfd, err := common.GetRktLockFD()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to get rkt lock fd: %v\n", err)
		return 1
	}

	if err := sys.CloseOnExec(lfd, true); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to set FD_CLOEXEC on rkt lock: %v\n", err)
		return 1
	}

	mirrorLocalZoneInfo(p.Root)

	if privNet {
		fps, err := forwardedPorts(p)
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			return 6
		}

		n, err := networking.Setup(root, p.UUID, fps)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to setup network: %v\n", err)
			return 6
		}
		defer n.Teardown()

		if err = n.Save(); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to save networking state %v\n", err)
			return 6
		}

		p.MetadataServiceURL = common.MetadataServicePublicURL(n.GetDefaultHostIP())

		if err = registerPod(p, n.GetDefaultIP()); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to register pod: %v\n", err)
			return 6
		}
		defer unregisterPod(p)
	}

	if err = p.PodToSystemd(interactive, virtualisation); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to configure systemd: %v\n", err)
		return 2
	}

	args, env, err := getArgsEnv(p)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to get execution parameters: %v\n", err)
		return 3
	}

	var execFn func() error

	if privNet {
		cmd := exec.Cmd{
			Path:   args[0],
			Args:   args,
			Stdin:  os.Stdin,
			Stdout: os.Stdout,
			Stderr: os.Stderr,
			Env:    env,
		}
		execFn = cmd.Run
	} else {
		execFn = func() error {
			return syscall.Exec(args[0], args, env)
		}
	}

	err = withClearedCloExec(lfd, execFn)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to execute containment: %v\n", err)
		return 5
	}

	return 0
}

func runningFromUnitFile() (ret bool, err error) {
	handle := C.dlopen(C.CString("libsystemd-login.so"), C.RTLD_LAZY)
	if handle == nil {
		// we can't open libsystemd-login.so so we assume systemd is not
		// installed and we're not running from a unit file
		ret = false
		return
	}
	defer func() {
		if r := C.dlclose(handle); r != 0 {
			err = fmt.Errorf("error closing libsystemd-login.so")
		}
	}()

	sd_pid_get_owner_uid := C.dlsym(handle, C.CString("sd_pid_get_owner_uid"))
	if sd_pid_get_owner_uid == nil {
		err = fmt.Errorf("error resolving sd_pid_get_owner_uid function")
		return
	}

	var uid C.uid_t
	errno := C.my_sd_pid_get_owner_uid(sd_pid_get_owner_uid, 0, &uid)
	// when we're running from a unit file, sd_pid_get_owner_uid returns
	// ENOENT (systemd <220) or ENXIO (systemd >=220)
	switch {
	case errno >= 0:
		ret = false
		return
	case syscall.Errno(-errno) == syscall.ENOENT || syscall.Errno(-errno) == syscall.ENXIO:
		ret = true
		return
	default:
		err = fmt.Errorf("error calling sd_pid_get_owner_uid: %v", syscall.Errno(-errno))
		return
	}
}

func main() {
	flag.Parse()

	if !debug {
		log.SetOutput(ioutil.Discard)
	}

	// move code into stage1() helper so defered fns get run
	os.Exit(stage1())
}
<commit_msg>Run lkvm in debug mode, when debugging is requested.<commit_after>// Copyright 2014 The rkt Authors
// Copyright 2015 Intel Corp
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//+build linux

package main

// #cgo LDFLAGS: -ldl
// #include <dlfcn.h>
// #include <sys/types.h>
//
// int
// my_sd_pid_get_owner_uid(void *f, pid_t pid, uid_t *uid)
// {
//   int (*sd_pid_get_owner_uid)(pid_t, uid_t *);
//
//   sd_pid_get_owner_uid = (int (*)(pid_t, uid_t *))f;
//   return sd_pid_get_owner_uid(pid, uid);
// }
//
import "C"

// this implements /init of stage1/nspawn+systemd

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"syscall"
	"strings"

	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/appc/spec/schema/types"

	"github.com/coreos/rkt/common"
	"github.com/coreos/rkt/networking"
	"github.com/coreos/rkt/pkg/sys"
)

const (
	// Path to systemd-nspawn binary within the stage1 rootfs
	nspawnBin = "/usr/bin/systemd-nspawn"
	// Path to lkvm binary within the stage1 rootfs
	lkvmBin = "/usr/bin/lkvm"
	bzImg = "/usr/lib/kernel/vmlinux.container"
	// Path to the interpreter within the stage1 rootfs
	interpBin = "/usr/lib64/ld-linux-x86-64.so.2"
	// Path to the localtime file/symlink in host
	localtimePath = "/etc/localtime"
)

// mirrorLocalZoneInfo tries to reproduce the /etc/localtime target in stage1/ to satisfy systemd-nspawn
func mirrorLocalZoneInfo(root string) {
	zif, err := os.Readlink(localtimePath)
	if err != nil {
		return
	}

	// On some systems /etc/localtime is a relative symlink, make it absolute
	if !filepath.IsAbs(zif) {
		zif = filepath.Join(filepath.Dir(localtimePath), zif)
		zif = filepath.Clean(zif)
	}

	src, err := os.Open(zif)
	if err != nil {
		return
	}
	defer src.Close()

	destp := filepath.Join(common.Stage1RootfsPath(root), zif)

	if err = os.MkdirAll(filepath.Dir(destp), 0755); err != nil {
		return
	}

	dest, err := os.OpenFile(destp, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer dest.Close()

	_, _ = io.Copy(dest, src)
}

var (
	debug       bool
	privNet     bool
	interactive bool
	virtualisation string
)

func init() {
	flag.BoolVar(&debug, "debug", false, "Run in debug mode")
	flag.BoolVar(&privNet, "private-net", false, "Setup private network")
	flag.BoolVar(&interactive, "interactive", false, "The pod is interactive")
	flag.StringVar(&virtualisation, "containment-type", "kvm", "Containment type to use: nspawn or kvm (default)")

	if os.Getenv("RKT_CONTAINMENT_TYPE") != "" {
		virtualisation = os.Getenv("RKT_CONTAINMENT_TYPE")
	}

	// this ensures that main runs only on main thread (thread group leader).
	// since namespace ops (unshare, setns) are done for a single thread, we
	// must ensure that the goroutine does not jump from OS thread to thread
	runtime.LockOSThread()
}

// getArgsEnvNspawn returns the nspawn args and env according to the usr used
func getArgsEnvNspawn(p *Pod) ([]string, []string, error) {
	args := []string{}
	env := os.Environ()

	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), interpBin))
	args = append(args, "--library-path")
	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), "usr/lib64"))

	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), nspawnBin))
	args = append(args, "--boot") // Launch systemd in the pod
	out, err := os.Getwd()
	if err != nil {
		return nil, nil, err
	}
	lfd, err := common.GetRktLockFD()
	if err != nil {
		return nil, nil, err
	}
	args = append(args, fmt.Sprintf("--pid-file=%v", filepath.Join(out, "pid")))
	args = append(args, fmt.Sprintf("--keep-fd=%v", lfd))
	args = append(args, fmt.Sprintf("--register=true"))

	if !debug {
		args = append(args, "--quiet") // silence most nspawn output (log_warning is currently not covered by this)
	}

	keepUnit, err := runningFromUnitFile()
	if err != nil {
		return nil, nil, fmt.Errorf("error determining if we're running from a unit file: %v", err)
	}

	if keepUnit {
		args = append(args, "--keep-unit")
	}

	nsargs, err := p.PodToNspawnArgs()
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to generate nspawn args: %v", err)
	}
	args = append(args, nsargs...)

	args = append(args, "--")
	args = append(args, "--default-standard-output=tty")

	if !debug {
		args = append(args, "--log-target=null")
		args = append(args, "--show-status=0")
	}

	return args, env, nil
}

func getArgsEnvKvm(p *Pod) ([]string, []string, error) {
	args := []string{}
	kargs := []string{}
	env := os.Environ()

	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), interpBin))
	args = append(args, "--library-path")
	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), "usr/lib64"))

	args = append(args, filepath.Join(common.Stage1RootfsPath(p.Root), lkvmBin))
	args = append(args, "run")

	args = append(args, "-m 1024")
	args = append(args, "-c 6")

	args = append(args, fmt.Sprintf("--kernel=%v", filepath.Join(common.Stage1RootfsPath(p.Root), bzImg)))
	args = append(args, "--console=virtio")
	kargs = append(kargs, "console=hvc0")

	kargs = append(kargs, "init=/usr/lib/systemd/systemd")

	nsargs, err := p.PodToKvmArgs()
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to generate kvm args: %v", err)
	}
	args = append(args, nsargs...)

	// Arguments to systemd
	kargs = append(kargs, "systemd.default_standard_output=tty")
	if debug {
		args = append(args, "--debug")
	} else {
		kargs = append(kargs, "systemd.log_target=null")
		kargs = append(kargs, "systemd.show-status=0")
		kargs = append(kargs, "quiet") // silence most nspawn output (log_warning is currently not covered by this)
	}

	args = append(args, "--param")
	args = append(args, strings.Join(kargs, " "))

	return args, env, nil
}

func getArgsEnv(p *Pod) ([]string, []string, error) {
	switch virtualisation {
	case "nspawn":
		return getArgsEnvNspawn(p)
	case "kvm":
		return getArgsEnvKvm(p)
	default:
		return nil, nil, fmt.Errorf("unrecognized containment type: %v", virtualisation)
	}
}

func withClearedCloExec(lfd int, f func() error) error {
	err := sys.CloseOnExec(lfd, false)
	if err != nil {
		return err
	}
	defer sys.CloseOnExec(lfd, true)

	return f()
}

func forwardedPorts(pod *Pod) ([]networking.ForwardedPort, error) {
	fps := []networking.ForwardedPort{}

	for _, ep := range pod.Manifest.Ports {
		n := ""
		fp := networking.ForwardedPort{}

		for _, a := range pod.Manifest.Apps {
			for _, p := range a.App.Ports {
				if p.Name == ep.Name {
					if n == "" {
						fp.Protocol = p.Protocol
						fp.HostPort = ep.HostPort
						fp.PodPort = p.Port
						n = a.Name.String()
					} else {
						return nil, fmt.Errorf("Ambiguous exposed port in PodManifest: %q and %q both define port %q", n, a.Name, p.Name)
					}
				}
			}
		}

		if n == "" {
			return nil, fmt.Errorf("Port name %q is not defined by any apps", ep.Name)
		}

		fps = append(fps, fp)
	}

	// TODO(eyakubovich): validate that there're no conflicts

	return fps, nil
}

func stage1() int {
	uuid, err := types.NewUUID(flag.Arg(0))
	if err != nil {
		fmt.Fprintln(os.Stderr, "UUID is missing or malformed")
		return 1
	}

	root := "."
	p, err := LoadPod(root, uuid)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load pod: %v\n", err)
		return 1
	}

	// set close-on-exec flag on RKT_LOCK_FD so it gets correctly closed when invoking
	// network plugins
	lfd, err := common.GetRktLockFD()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to get rkt lock fd: %v\n", err)
		return 1
	}

	if err := sys.CloseOnExec(lfd, true); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to set FD_CLOEXEC on rkt lock: %v\n", err)
		return 1
	}

	mirrorLocalZoneInfo(p.Root)

	if privNet {
		fps, err := forwardedPorts(p)
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			return 6
		}

		n, err := networking.Setup(root, p.UUID, fps)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to setup network: %v\n", err)
			return 6
		}
		defer n.Teardown()

		if err = n.Save(); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to save networking state %v\n", err)
			return 6
		}

		p.MetadataServiceURL = common.MetadataServicePublicURL(n.GetDefaultHostIP())

		if err = registerPod(p, n.GetDefaultIP()); err != nil {
			fmt.Fprintf(os.Stderr, "Failed to register pod: %v\n", err)
			return 6
		}
		defer unregisterPod(p)
	}

	if err = p.PodToSystemd(interactive, virtualisation); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to configure systemd: %v\n", err)
		return 2
	}

	args, env, err := getArgsEnv(p)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to get execution parameters: %v\n", err)
		return 3
	}

	var execFn func() error

	if privNet {
		cmd := exec.Cmd{
			Path:   args[0],
			Args:   args,
			Stdin:  os.Stdin,
			Stdout: os.Stdout,
			Stderr: os.Stderr,
			Env:    env,
		}
		execFn = cmd.Run
	} else {
		execFn = func() error {
			return syscall.Exec(args[0], args, env)
		}
	}

	err = withClearedCloExec(lfd, execFn)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to execute containment: %v\n", err)
		return 5
	}

	return 0
}

func runningFromUnitFile() (ret bool, err error) {
	handle := C.dlopen(C.CString("libsystemd-login.so"), C.RTLD_LAZY)
	if handle == nil {
		// we can't open libsystemd-login.so so we assume systemd is not
		// installed and we're not running from a unit file
		ret = false
		return
	}
	defer func() {
		if r := C.dlclose(handle); r != 0 {
			err = fmt.Errorf("error closing libsystemd-login.so")
		}
	}()

	sd_pid_get_owner_uid := C.dlsym(handle, C.CString("sd_pid_get_owner_uid"))
	if sd_pid_get_owner_uid == nil {
		err = fmt.Errorf("error resolving sd_pid_get_owner_uid function")
		return
	}

	var uid C.uid_t
	errno := C.my_sd_pid_get_owner_uid(sd_pid_get_owner_uid, 0, &uid)
	// when we're running from a unit file, sd_pid_get_owner_uid returns
	// ENOENT (systemd <220) or ENXIO (systemd >=220)
	switch {
	case errno >= 0:
		ret = false
		return
	case syscall.Errno(-errno) == syscall.ENOENT || syscall.Errno(-errno) == syscall.ENXIO:
		ret = true
		return
	default:
		err = fmt.Errorf("error calling sd_pid_get_owner_uid: %v", syscall.Errno(-errno))
		return
	}
}

func main() {
	flag.Parse()

	if !debug {
		log.SetOutput(ioutil.Discard)
	}

	// move code into stage1() helper so defered fns get run
	os.Exit(stage1())
}
<|endoftext|>


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

//Sanity Check Packet -> minimum length and correct checksum
//decrement TTL and recompute packet checksum (l3 recompute checksum)

//lookup in the longest prefix matching table:
//destination ip address of the packet.

//LONGEST PREFIX MATCHING trivialimplementation

var RouterCode = `
#include <linux/ip.h>
#include <linux/bpf.h>
#include <linux/kernel.h>

// #define BPF_TRACE
#undef BPF_TRACE

#define BPF_LOG
// #undef BPF_LOG

#define ROUTING_TABLE_DIM 10
#define ROUTER_PORT_N     10
#define ARP_TABLE_DIM     10

#define IP_TTL_OFFSET  8
#define IP_CSUM_OFFSET 10

#define ETH_DST_OFFSET  0
#define ETH_SRC_OFFSET  6
#define ETH_TYPE_OFFSET 12

/*Routing Table Entry*/
struct rt_entry{
  u32 network;  //network: e.g. 192.168.1.0
  u32 netmask;  //netmask: e.g. 255.255.255.0
  u32 port;     //port of the router
};

/*Router Port*/
struct r_port{
  u32 ip;       //ip addr : e.g. 192.168.1.254
  u32 netmask;  //netmask : e.g. 255.255.255.0
  u64 mac;      //mac addr: e.g. a1:b2:c3:ab:cd:ef
};

/*Arp Table Key*/
struct arp_table_key{
  u32 ip;       //ip addr : e.g. 192.168.1.2
  u32 port;     //port    : e.g. 1
};

/*
  The Routing table is implemented as an array of struct rt_entry (Routing Table Entry)
  the longest prefix matching algorithm (at least a simplified version)
  is implemented performing a bounded loop over the entries of the routing table.
  We assume that the control plane puts entry ordered from the longest netmask
  to the shortest one.
*/
BPF_TABLE("array", u32, struct rt_entry, routing_table, ROUTING_TABLE_DIM);

/*
  Router Port table provides a way to simulate the physical interface of the router
  The ip address is used to answer to the arp request (TO IMPLEMENT)
  The mac address is used as mac_scr for the outcoming packet on that interface,
  and as mac address contained in the arp reply
*/
BPF_TABLE("hash", u32, struct r_port, router_port, ROUTER_PORT_N);

/*
  We shold have an arp table for each port of the router?
  For now we assume to send packet exiting the router interfaces in broadcast
  (mac dst = ff:ff:ff:ff:ff:ff)

  How can we implement multiple arp tables?
  One possible implementation using one single map is the following
  key{ ip + port number } -> value {mac_address}
*/
BPF_TABLE("hash", u32, u64, arp_table, ARP_TABLE_DIM);

static int handle_rx(void *skb, struct metadata *md) {
  u8 *cursor = 0;
  struct ethernet_t *ethernet = cursor_advance(cursor, sizeof(*ethernet));

  #ifdef BPF_TRACE
    bpf_trace_printk("[router-%d]: in_ifc:%d\n", md->module_id, md->in_ifc);
    bpf_trace_printk("[router-%d]: eth_type:%x mac_scr:%lx mac_dst:%lx\n",
      md->module_id, ethernet->type, ethernet->src, ethernet->dst);
  #endif

  //TODO
  //sanity check of the packet.
  //if something wrong -> DROP the packet

  // is it an ipv4 packet?
  if (ethernet->type == 0x0800) {
    struct ip_t *ip = cursor_advance(cursor, sizeof(*ip));

    #ifdef BPF_TRACE
      bpf_trace_printk("[router-%d]: ttl:%u ip_scr:%x ip_dst:%x \n", md->module_id, ip->ttl, ip->src, ip->dst);
      // bpf_trace_printk("[router-%d]: (before) ttl: %d checksum: %x\n", ip->ttl, ip->hchecksum);
    #endif

    /*
      decrement TTL and recompute packet checksum (l3 recompute checksum).
      if ttl <= 1 DROP the packet.
      eventually send ICMP message for the packet dropped.
      (maybe to avoid for security reasons)
    */

    __u8 old_ttl = ip->ttl;
    __u8 new_ttl;

    if (old_ttl <= 1) {
      #ifdef BPF_TRACE
        bpf_trace_printk("[router-%d]: packet DROP (ttl <= 1)\n", md->module_id);
      #endif
      return RX_DROP;
    }

    new_ttl = old_ttl - 1;
    bpf_l3_csum_replace(skb, sizeof(*ethernet) + IP_CSUM_OFFSET , old_ttl, new_ttl, sizeof(__u16));
    bpf_skb_store_bytes(skb, sizeof(*ethernet) + IP_TTL_OFFSET , &new_ttl, sizeof(old_ttl), 0);

    #ifdef BPF_TRACE
      // bpf_trace_printk("[router-%d]: (after ) ttl: %d checksum: %x\n",ip->ttl,ip->hchecksum);
    #endif

    /*
      ROUTING ALGORITHM (simplified)

      for each item in the routing table (upbounded loop)
      apply the netmask on dst_ip_address
      (possible optimization, not recompute if at next iteration the netmask is the same)
      if masked address == network in the routing table
        1- change src mac to otuput port mac
        2- change dst mac to lookup arp table (or send to fffffffffffff)
        3- forward the packet to dst port
    */

    int i = 0;
    struct rt_entry *rt_entry_p = 0;

    u64 new_src_mac = 0;
    u64 new_dst_mac = 0;
    u32 out_port = 0;
    struct r_port *r_port_p = 0;

    #pragma unroll
    for (i = 0; i < ROUTING_TABLE_DIM; i++) {
      u32 t = i;
      rt_entry_p = routing_table.lookup(&t);
       if (rt_entry_p) {
        if ((ip->dst & rt_entry_p->netmask) == rt_entry_p->network) {
          goto FORWARD;
        }
      }
    }

  DROP:
    #ifdef BPF_LOG
      bpf_trace_printk("[router-%d]: in: %d out: -- DROP\n", md->module_id, md->in_ifc);
    #endif
    return RX_DROP;

  FORWARD:
    //Select out interface
    out_port = rt_entry_p->port;
    if (out_port <= 0)
      goto DROP;

    #ifdef BPF_LOG
      bpf_trace_printk("[router-%d]: routing table match (#%d) network: %x\n",
        md->module_id, i, rt_entry_p->network);
    #endif

    //change src mac
    r_port_p = router_port.lookup(&out_port);
    if (r_port_p) {
      new_src_mac = cpu_to_be64(r_port_p->mac<<16);
      bpf_skb_store_bytes(skb,ETH_SRC_OFFSET, &new_src_mac, 6, 0);
    }

    //change dst mac to ff:ff:ff:ff:ff:ff (TODO arp table)
    new_dst_mac = 0xffffffffffff;
    bpf_skb_store_bytes(skb, ETH_DST_OFFSET, &new_dst_mac, 6, 0);

    #ifdef BPF_TRACE
      bpf_trace_printk("[router-%d]: eth_type:%x mac_scr:%lx mac_dst:%lx\n",
        md->module_id, ethernet->type, ethernet->src, ethernet->dst);
      bpf_trace_printk("[router-%d]: out_ifc: %d\n", out_port);
    #endif

    #ifdef BPF_LOG
      bpf_trace_printk("[router-%d]: in: %d out: %d REDIRECT\n", md->module_id, md->in_ifc, out_port);
    #endif

    pkt_redirect(skb,md,out_port);
    return RX_REDIRECT;
  }
  else if(ethernet->type == 0x0806) { // is it ARP?
    struct arp_t *arp = cursor_advance(cursor, sizeof(*arp));
    if (arp->oper == 1) {	// arp request?
      //bpf_trace_printk("[arp]: packet is arp request\n");

      struct r_port *port = router_port.lookup(&md->in_ifc);
      if (!port)
        return RX_DROP;
      if (arp->tpa == port->ip) {
        //bpf_trace_printk("[arp]: Somebody is asking for my address\n");

        /* due to a bcc issue: https://github.com/iovisor/bcc/issues/537 it
         * is necessary to copy the data field into a temporal variable
         */
        u64 mymac = port->mac;
        u64 remotemac = arp->sha;
        u32 myip = port->ip;
        u32 remoteip = arp->spa;

        ethernet->dst = remotemac;
        ethernet->src = mymac;

        /* please note that the mac has to be copied before that the ips.  This
         * is because the temporal variable used to save the mac has 8 byes, 2
         * more than the mac itself.  Then when copying the mac into the packet
         * the two first bytes of the ip are also modified.
         */
        arp->oper = 2;
        arp->tha = remotemac;
        arp->sha = mymac;
        arp->tpa = remoteip;
        arp->spa = myip;

        /* register the requesting mac and ips */
        arp_table.update(&remoteip, &remotemac);

        /* register the requesting mac and ips */
        arp_table.update(&remoteip, &remotemac);

        pkt_redirect(skb, md, md->in_ifc);

        return RX_REDIRECT;

      }
    }
    else if (arp->oper == 2) { //arp reply
      bpf_trace_printk("[router-%d]: packet is arp reply\n", md->module_id);

      struct r_port *port = router_port.lookup(&md->in_ifc);
      if (!port)
        return RX_DROP;
      if (arp->sha == port->mac && arp->spa == port->ip) {
        u64 mac_ = port->mac;
        u32 ip_ = port->ip;
        arp_table.update(&ip_, &mac_);
        return RX_DROP;
      }
    }
  }

  return RX_DROP;
}
`
<commit_msg>iomodules/router: change mac destination based on arp table<commit_after>// Copyright 2016 Politecnico di Torino
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package router

//Sanity Check Packet -> minimum length and correct checksum
//decrement TTL and recompute packet checksum (l3 recompute checksum)

//lookup in the longest prefix matching table:
//destination ip address of the packet.

//LONGEST PREFIX MATCHING trivialimplementation

var RouterCode = `
#include <linux/ip.h>
#include <linux/bpf.h>
#include <linux/kernel.h>

// #define BPF_TRACE
#undef BPF_TRACE

#define BPF_LOG
// #undef BPF_LOG

#define ROUTING_TABLE_DIM 10
#define ROUTER_PORT_N     10
#define ARP_TABLE_DIM     10

#define IP_TTL_OFFSET  8
#define IP_CSUM_OFFSET 10

#define ETH_DST_OFFSET  0
#define ETH_SRC_OFFSET  6
#define ETH_TYPE_OFFSET 12

/*Routing Table Entry*/
struct rt_entry{
  u32 network;  //network: e.g. 192.168.1.0
  u32 netmask;  //netmask: e.g. 255.255.255.0
  u32 port;     //port of the router
};

/*Router Port*/
struct r_port{
  u32 ip;       //ip addr : e.g. 192.168.1.254
  u32 netmask;  //netmask : e.g. 255.255.255.0
  u64 mac;      //mac addr: e.g. a1:b2:c3:ab:cd:ef
};

/*Arp Table Key*/
struct arp_table_key{
  u32 ip;       //ip addr : e.g. 192.168.1.2
  u32 port;     //port    : e.g. 1
};

/*
  The Routing table is implemented as an array of struct rt_entry (Routing Table Entry)
  the longest prefix matching algorithm (at least a simplified version)
  is implemented performing a bounded loop over the entries of the routing table.
  We assume that the control plane puts entry ordered from the longest netmask
  to the shortest one.
*/
BPF_TABLE("array", u32, struct rt_entry, routing_table, ROUTING_TABLE_DIM);

/*
  Router Port table provides a way to simulate the physical interface of the router
  The ip address is used to answer to the arp request (TO IMPLEMENT)
  The mac address is used as mac_scr for the outcoming packet on that interface,
  and as mac address contained in the arp reply
*/
BPF_TABLE("hash", u32, struct r_port, router_port, ROUTER_PORT_N);

/*
  We shold have an arp table for each port of the router?
  For now we assume to send packet exiting the router interfaces in broadcast
  (mac dst = ff:ff:ff:ff:ff:ff)

  How can we implement multiple arp tables?
  One possible implementation using one single map is the following
  key{ ip + port number } -> value {mac_address}
*/
BPF_TABLE("hash", u32, u64, arp_table, ARP_TABLE_DIM);

static int handle_rx(void *skb, struct metadata *md) {
  u8 *cursor = 0;
  struct ethernet_t *ethernet = cursor_advance(cursor, sizeof(*ethernet));

  #ifdef BPF_TRACE
    bpf_trace_printk("[router-%d]: in_ifc:%d\n", md->module_id, md->in_ifc);
    bpf_trace_printk("[router-%d]: eth_type:%x mac_scr:%lx mac_dst:%lx\n",
      md->module_id, ethernet->type, ethernet->src, ethernet->dst);
  #endif

  //TODO
  //sanity check of the packet.
  //if something wrong -> DROP the packet

  // is it an ipv4 packet?
  if (ethernet->type == 0x0800) {
    struct ip_t *ip = cursor_advance(cursor, sizeof(*ip));

    #ifdef BPF_TRACE
      bpf_trace_printk("[router-%d]: ttl:%u ip_scr:%x ip_dst:%x \n", md->module_id, ip->ttl, ip->src, ip->dst);
      // bpf_trace_printk("[router-%d]: (before) ttl: %d checksum: %x\n", ip->ttl, ip->hchecksum);
    #endif

    /*
      decrement TTL and recompute packet checksum (l3 recompute checksum).
      if ttl <= 1 DROP the packet.
      eventually send ICMP message for the packet dropped.
      (maybe to avoid for security reasons)
    */

    __u8 old_ttl = ip->ttl;
    __u8 new_ttl;

    if (old_ttl <= 1) {
      #ifdef BPF_TRACE
        bpf_trace_printk("[router-%d]: packet DROP (ttl <= 1)\n", md->module_id);
      #endif
      return RX_DROP;
    }

    new_ttl = old_ttl - 1;
    bpf_l3_csum_replace(skb, sizeof(*ethernet) + IP_CSUM_OFFSET , old_ttl, new_ttl, sizeof(__u16));
    bpf_skb_store_bytes(skb, sizeof(*ethernet) + IP_TTL_OFFSET , &new_ttl, sizeof(old_ttl), 0);

    #ifdef BPF_TRACE
      // bpf_trace_printk("[router-%d]: (after ) ttl: %d checksum: %x\n",ip->ttl,ip->hchecksum);
    #endif

    /*
      ROUTING ALGORITHM (simplified)

      for each item in the routing table (upbounded loop)
      apply the netmask on dst_ip_address
      (possible optimization, not recompute if at next iteration the netmask is the same)
      if masked address == network in the routing table
        1- change src mac to otuput port mac
        2- change dst mac to lookup arp table (or send to fffffffffffff)
        3- forward the packet to dst port
    */

    int i = 0;
    struct rt_entry *rt_entry_p = 0;

    u32 out_port = 0;
    struct r_port *r_port_p = 0;

    #pragma unroll
    for (i = 0; i < ROUTING_TABLE_DIM; i++) {
      u32 t = i;
      rt_entry_p = routing_table.lookup(&t);
       if (rt_entry_p) {
        if ((ip->dst & rt_entry_p->netmask) == rt_entry_p->network) {
          goto FORWARD;
        }
      }
    }

  DROP:
    #ifdef BPF_LOG
      bpf_trace_printk("[router-%d]: in: %d out: -- DROP\n", md->module_id, md->in_ifc);
    #endif
    return RX_DROP;

  FORWARD:
    //Select out interface
    out_port = rt_entry_p->port;
    if (out_port <= 0)
      goto DROP;

    #ifdef BPF_LOG
      bpf_trace_printk("[router-%d]: routing table match (#%d) network: %x\n",
        md->module_id, i, rt_entry_p->network);
    #endif

    //change src mac
    r_port_p = router_port.lookup(&out_port);
    if (r_port_p) {
      ethernet->src = r_port_p->mac;
    }

    //change dst mac
    u32 dst_ip = ip->dst;
    u64 new_dst_mac = 0xffffffffffff;
    u64 *mac_entry = arp_table.lookup(&dst_ip);
    if (mac_entry) {
      new_dst_mac = *mac_entry;
    }

    ethernet->dst = new_dst_mac;

    #ifdef BPF_TRACE
      bpf_trace_printk("[router-%d]: eth_type:%x mac_scr:%lx mac_dst:%lx\n",
        md->module_id, ethernet->type, ethernet->src, ethernet->dst);
      bpf_trace_printk("[router-%d]: out_ifc: %d\n", out_port);
    #endif

    #ifdef BPF_LOG
      bpf_trace_printk("[router-%d]: in: %d out: %d REDIRECT\n", md->module_id, md->in_ifc, out_port);
    #endif

    pkt_redirect(skb,md,out_port);
    return RX_REDIRECT;
  }
  else if(ethernet->type == 0x0806) { // is it ARP?
    struct arp_t *arp = cursor_advance(cursor, sizeof(*arp));
    if (arp->oper == 1) {	// arp request?
      //bpf_trace_printk("[arp]: packet is arp request\n");

      struct r_port *port = router_port.lookup(&md->in_ifc);
      if (!port)
        return RX_DROP;
      if (arp->tpa == port->ip) {
        //bpf_trace_printk("[arp]: Somebody is asking for my address\n");

        /* due to a bcc issue: https://github.com/iovisor/bcc/issues/537 it
         * is necessary to copy the data field into a temporal variable
         */
        u64 mymac = port->mac;
        u64 remotemac = arp->sha;
        u32 myip = port->ip;
        u32 remoteip = arp->spa;

        ethernet->dst = remotemac;
        ethernet->src = mymac;

        /* please note that the mac has to be copied before that the ips.  This
         * is because the temporal variable used to save the mac has 8 byes, 2
         * more than the mac itself.  Then when copying the mac into the packet
         * the two first bytes of the ip are also modified.
         */
        arp->oper = 2;
        arp->tha = remotemac;
        arp->sha = mymac;
        arp->tpa = remoteip;
        arp->spa = myip;

        /* register the requesting mac and ips */
        arp_table.update(&remoteip, &remotemac);

        /* register the requesting mac and ips */
        arp_table.update(&remoteip, &remotemac);

        pkt_redirect(skb, md, md->in_ifc);

        return RX_REDIRECT;

      }
    }
    else if (arp->oper == 2) { //arp reply
      bpf_trace_printk("[router-%d]: packet is arp reply\n", md->module_id);

      struct r_port *port = router_port.lookup(&md->in_ifc);
      if (!port)
        return RX_DROP;
      if (arp->sha == port->mac && arp->spa == port->ip) {
        u64 mac_ = port->mac;
        u32 ip_ = port->ip;
        arp_table.update(&ip_, &mac_);
        return RX_DROP;
      }
    }
  }

  return RX_DROP;
}
`
<|endoftext|>