<commit_before>// bootnode runs a bootstrap node for the Ethereum Discovery Protocol.
package main

import (
	"crypto/ecdsa"
	"encoding/hex"
	"flag"
	"fmt"
	"net"
	"os"

	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/log"
	"github.com/ethereum/go-ethereum/p2p/discv5"
)

type bootnodes []*discv5.Node

func (f *bootnodes) String() string {
	return "discv5 nodes"
}

// Set unmarshals enode into discv5.Node.
func (f *bootnodes) Set(value string) error {
	n, err := discv5.ParseNode(value)
	if err != nil {
		return err
	}
	*f = append(*f, n)
	return nil
}

func main() {
	var (
		writeAddr   = flag.Bool("writeaddress", false, "write out the node's public key and quit")
		listenAddr  = flag.String("addr", ":30301", "listen address")
		genKeyFile  = flag.String("genkey", "", "generate a node key")
		nodeKeyFile = flag.String("nodekey", "", "private key filename")
		keydata     = flag.String("keydata", "", "hex encoded private key")
		verbosity   = flag.Int("verbosity", int(log.LvlInfo), "log verbosity (0-9)")
		vmodule     = flag.String("vmodule", "", "log verbosity pattern")
		nursery     = bootnodes{}
		nodeKey     *ecdsa.PrivateKey
		err         error
	)
	flag.Var(&nursery, "n", "These nodes are used to connect to the network if the table is empty and there are no known nodes in the database.")
	flag.Parse()

	glogger := log.NewGlogHandler(log.StreamHandler(os.Stderr, log.TerminalFormat(false)))
	glogger.Verbosity(log.Lvl(*verbosity))
	if err = glogger.Vmodule(*vmodule); err != nil {
		log.Crit("Failed to set glog verbosity", "value", *vmodule, "err", err)
	}
	log.Root().SetHandler(glogger)

	if len(*genKeyFile) != 0 {
		log.Info("Generating key file", "path", *genKeyFile)
		key, err := crypto.GenerateKey()
		if err != nil {
			log.Crit("unable to generate key", "error", err)
		}
		if err := crypto.SaveECDSA(*genKeyFile, key); err != nil {
			log.Crit("unable to save key", "error", err)
		}
		os.Exit(0)
	}
	if len(*nodeKeyFile) == 0 && len(*keydata) == 0 {
		log.Crit("either `nodekey` or `keydata` must be provided")
	}
	if len(*nodeKeyFile) != 0 {
		nodeKey, err = crypto.LoadECDSA(*nodeKeyFile)
		if err != nil {
			log.Crit("Failed to load ecdsa key from", "file", *nodeKeyFile, "error", err)
		}
	} else if len(*keydata) != 0 {
		log.Warn("key will be visible in process list. should be used only for tests")
		key, err := hex.DecodeString(*keydata)
		if err != nil {
			log.Crit("unable to decode hex", "data", keydata, "error", err)
		}
		nodeKey, err = crypto.ToECDSA(key)
		if err != nil {
			log.Crit("unable to convert decoded hex into ecdsa.PrivateKey", "data", key, "error", err)
		}
	}
	if *writeAddr {
		// we remove the first uncompressed byte since it's not used in an enode address
		fmt.Printf("%x\n", crypto.FromECDSAPub(&nodeKey.PublicKey)[1:])
		os.Exit(0)
	}

	addr, err := net.ResolveUDPAddr("udp", *listenAddr)
	if err != nil {
		log.Crit("Unable to resolve UDP", "address", *listenAddr, "error", err)
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Crit("Unable to listen on udp", "address", addr, "error", err)
	}

	tab, err := discv5.ListenUDP(nodeKey, conn, "", nil)
	if err != nil {
		log.Crit("Failed to create discovery v5 table:", "error", err)
	}
	defer tab.Close()
	if err := tab.SetFallbackNodes(nursery); err != nil {
		log.Crit("Failed to set fallback", "nodes", nursery, "error", err)
	}
	select {}
}
<commit_msg>fix linter gocyclo error for bootnode<commit_after>// bootnode runs a bootstrap node for the Ethereum Discovery Protocol.
package main

import (
	"crypto/ecdsa"
	"encoding/hex"
	"flag"
	"fmt"
	"net"
	"os"

	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/log"
	"github.com/ethereum/go-ethereum/p2p/discv5"
)

var (
	writeAddr   = flag.Bool("writeaddress", false, "write out the node's public key and quit")
	listenAddr  = flag.String("addr", ":30301", "listen address")
	genKeyFile  = flag.String("genkey", "", "generate a node key")
	nodeKeyFile = flag.String("nodekey", "", "private key filename")
	keydata     = flag.String("keydata", "", "hex encoded private key")
	verbosity   = flag.Int("verbosity", int(log.LvlInfo), "log verbosity (0-9)")
	vmodule     = flag.String("vmodule", "", "log verbosity pattern")
	nursery     = bootnodes{}
	nodeKey     *ecdsa.PrivateKey
	err         error
)

type bootnodes []*discv5.Node

func (f *bootnodes) String() string {
	return "discv5 nodes"
}

// Set unmarshals enode into discv5.Node.
func (f *bootnodes) Set(value string) error {
	n, err := discv5.ParseNode(value)
	if err != nil {
		return err
	}
	*f = append(*f, n)
	return nil
}

func main() { // nolint: gocyclo
	flag.Var(&nursery, "n", "These nodes are used to connect to the network if the table is empty and there are no known nodes in the database.")
	flag.Parse()

	glogger := log.NewGlogHandler(log.StreamHandler(os.Stderr, log.TerminalFormat(false)))
	glogger.Verbosity(log.Lvl(*verbosity))
	if err = glogger.Vmodule(*vmodule); err != nil {
		log.Crit("Failed to set glog verbosity", "value", *vmodule, "err", err)
	}
	log.Root().SetHandler(glogger)

	if len(*genKeyFile) != 0 {
		log.Info("Generating key file", "path", *genKeyFile)
		key, err := crypto.GenerateKey()
		if err != nil {
			log.Crit("unable to generate key", "error", err)
		}
		if err := crypto.SaveECDSA(*genKeyFile, key); err != nil {
			log.Crit("unable to save key", "error", err)
		}
		os.Exit(0)
	}
	if len(*nodeKeyFile) == 0 && len(*keydata) == 0 {
		log.Crit("either `nodekey` or `keydata` must be provided")
	}
	if len(*nodeKeyFile) != 0 {
		nodeKey, err = crypto.LoadECDSA(*nodeKeyFile)
		if err != nil {
			log.Crit("Failed to load ecdsa key from", "file", *nodeKeyFile, "error", err)
		}
	} else if len(*keydata) != 0 {
		log.Warn("key will be visible in process list. should be used only for tests")
		key, err := hex.DecodeString(*keydata)
		if err != nil {
			log.Crit("unable to decode hex", "data", keydata, "error", err)
		}
		nodeKey, err = crypto.ToECDSA(key)
		if err != nil {
			log.Crit("unable to convert decoded hex into ecdsa.PrivateKey", "data", key, "error", err)
		}
	}
	if *writeAddr {
		// we remove the first uncompressed byte since it's not used in an enode address
		fmt.Printf("%x\n", crypto.FromECDSAPub(&nodeKey.PublicKey)[1:])
		os.Exit(0)
	}

	addr, err := net.ResolveUDPAddr("udp", *listenAddr)
	if err != nil {
		log.Crit("Unable to resolve UDP", "address", *listenAddr, "error", err)
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Crit("Unable to listen on udp", "address", addr, "error", err)
	}

	tab, err := discv5.ListenUDP(nodeKey, conn, "", nil)
	if err != nil {
		log.Crit("Failed to create discovery v5 table:", "error", err)
	}
	defer tab.Close()
	if err := tab.SetFallbackNodes(nursery); err != nil {
		log.Crit("Failed to set fallback", "nodes", nursery, "error", err)
	}

	select {}
}
<|endoftext|>


<commit_before>package main

import (
    "os"
    "os/exec"
    "bufio"
    "fmt"
    "flag"
    "encoding/json"
    "io/ioutil"
    "github.com/speedata/gogit"
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3/s3manager"
    "github.com/deckarep/golang-set"
    "path/filepath"
)

const REF_S3_PUSH string = "LAST_S3_PUSH"
const CONFIG_FILE_PATH string = ".git_s3_push"

type Repository struct {
    GitRepo         *gogit.Repository
    HeadCommit      *gogit.Commit
    LastPushCommit  *gogit.Commit
    UnpushedFiles   mapset.Set
    Config          RepoConfig
    S3Uploader      S3Uploader
}

type RepoConfig struct {
    S3Region        string
    S3Bucket        string
}

func OpenRepository() (*Repository, error) {
    repo := new(Repository)
    repo.UnpushedFiles = mapset.NewSet()

    wd, err := os.Getwd()
    if err != nil {
        return nil, err
    }

    path := filepath.Join(wd, ".git")
    if _, err := os.Stat(path); os.IsNotExist(err) {
        return nil, err
    }

    gitRepo, err := gogit.OpenRepository(path)
    if err != nil {
        return nil, err
    }
    repo.GitRepo = gitRepo

    return repo, nil
}

func (repo *Repository) ReadConfigFile() error {
    file, err := ioutil.ReadFile(CONFIG_FILE_PATH)
    if err != nil {
        return err
    }

    err = json.Unmarshal(file, &repo.Config)
    if err != nil {
        return err
    }
    return nil
}

func (repo Repository) SaveConfigToFile() error {
    jsonData, err := json.Marshal(repo.Config)
    if err != nil {
        return err
    }

    err = ioutil.WriteFile(CONFIG_FILE_PATH, jsonData, 0644)
    if err != nil {
        return err
    }

    return nil
}

func (repo *Repository) FindRelevantCommits() error {
    headRef, err := repo.GitRepo.LookupReference("HEAD")
    if err != nil {
        return err
    }

    headCommit, err := repo.GitRepo.LookupCommit(headRef.Target())
    if err != nil {
        return err
    }
    repo.HeadCommit = headCommit

    lastPushRef, err := repo.GitRepo.LookupReference(REF_S3_PUSH)
    if err != nil {
        return nil
    }

    lastPushCommit, err := repo.GitRepo.LookupCommit(lastPushRef.Target())
    if err != nil {
        return nil
    }
    repo.LastPushCommit = lastPushCommit

    return nil
}

func (repo *Repository) ReadGitModifiedFiles(scanner *bufio.Scanner, stop chan bool)  {
    for scanner.Scan() {
        repo.UnpushedFiles.Add(scanner.Text())
    }

    stop <- true
}

func (repo *Repository) FindCommitModifiedFiles(commit *gogit.Commit) error {
    cmd := exec.Command("git", "diff-tree", "--no-commit-id", "--name-only", "--root", commit.Id().String())
    out, err := cmd.StdoutPipe()
    if err != nil {
        return err
    }

    err = cmd.Start()
    if err != nil {
        return err
    }

    scanner := bufio.NewScanner(out)

    stop := make(chan bool)
    go repo.ReadGitModifiedFiles(scanner, stop)
    <-stop
    cmd.Wait()

    return nil
}

func (repo *Repository) FindUnpushedModifiedFiles() error {
    queue := []*gogit.Commit{};
    visited := mapset.NewSet();

    currentCommit := repo.HeadCommit;
    for currentCommit != nil && currentCommit.ParentCount() > 0 {
        if repo.LastPushCommit != nil && repo.LastPushCommit.Id().Equal(currentCommit.Id()) {
            break;
        }

        err := repo.FindCommitModifiedFiles(currentCommit)
        if err != nil {
            return err
        }

        for i := 0; i < currentCommit.ParentCount(); i++ {
            parentCommit := currentCommit.Parent(i)
            if !visited.Contains(parentCommit) {
                queue = append(queue, parentCommit)
            }
        }

        if len(queue) < 1 {
            break;
        }

        currentCommit = queue[0]
        queue = queue[1:]
    }

    return nil
}

type S3Uploader struct {
    BucketName      *string
    S3Uploader      *s3manager.Uploader
}

func InitS3Uploader(config RepoConfig) *S3Uploader {
    uploader := new(S3Uploader)
    uploader.BucketName = aws.String(config.S3Bucket)

    s3config := aws.Config{Region: aws.String(config.S3Region)}
    s3uploader := s3manager.NewUploader(session.New(&s3config))
    uploader.S3Uploader = s3uploader

    return uploader
}

func (uploader S3Uploader) UploadFile(path string) error {
    file, err := os.Open(path)
    if err != nil {
        return err
    }

    result, err := uploader.S3Uploader.Upload(&s3manager.UploadInput{
        Body: file,
        Bucket: uploader.BucketName,
        Key: aws.String(path),
    })

    if err != nil {
        return err
    }

    fmt.Println(result.Location)
    return nil
}

func main() {
    repo, err := OpenRepository()
    if err != nil {
        fmt.Println(err)
        os.Exit(1)
    }
    repo.ReadConfigFile()

    flag.StringVar(&repo.Config.S3Bucket, "b", repo.Config.S3Bucket, "Destination S3 bucket name")
    flag.StringVar(&repo.Config.S3Region, "r", repo.Config.S3Region, "AWS region of destination bucket")
    saveConfig := flag.Bool("save", false, "Save destination region/bucket to config file")
    flag.Parse()

    if repo.Config.S3Bucket == "" {
        flag.Usage()
        os.Exit(1)
    } else if (repo.Config.S3Region == "") {
        flag.Usage()
        os.Exit(1)
    } else if (*saveConfig) {
        err := repo.SaveConfigToFile()
        if err != nil {
            fmt.Println("WARNING: Failed to save config to file: ", err)
        }
    }

    if err := repo.FindRelevantCommits(); err != nil {
        fmt.Println(err)
        os.Exit(1)
    }

    repo.FindUnpushedModifiedFiles();

    if repo.UnpushedFiles.Cardinality() == 0 {
        fmt.Println("No modified files to push")
        os.Exit(0)
    }

    uploader := InitS3Uploader(repo.Config)

    for filePath := range repo.UnpushedFiles.Iter() {
        fmt.Println("Uploading: " + filePath.(string))
        err := uploader.UploadFile(filePath.(string))
        if err != nil {
            fmt.Println(err)
            os.Exit(1)
        }
    }
}<commit_msg>Update LAST_S3_PUSH ref in git after successful push to S3<commit_after>package main

import (
    "os"
    "os/exec"
    "bufio"
    "fmt"
    "flag"
    "encoding/json"
    "io/ioutil"
    "github.com/speedata/gogit"
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3/s3manager"
    "github.com/deckarep/golang-set"
    "path/filepath"
)

const REF_S3_PUSH string = "LAST_S3_PUSH"
const CONFIG_FILE_PATH string = ".git_s3_push"

type Repository struct {
    GitRepo         *gogit.Repository
    HeadCommit      *gogit.Commit
    LastPushCommit  *gogit.Commit
    UnpushedFiles   mapset.Set
    Config          RepoConfig
    S3Uploader      S3Uploader
}

type RepoConfig struct {
    S3Region        string
    S3Bucket        string
}

func OpenRepository() (*Repository, error) {
    repo := new(Repository)
    repo.UnpushedFiles = mapset.NewSet()

    wd, err := os.Getwd()
    if err != nil {
        return nil, err
    }

    path := filepath.Join(wd, ".git")
    if _, err := os.Stat(path); os.IsNotExist(err) {
        return nil, err
    }

    gitRepo, err := gogit.OpenRepository(path)
    if err != nil {
        return nil, err
    }
    repo.GitRepo = gitRepo

    return repo, nil
}

func (repo *Repository) ReadConfigFile() error {
    file, err := ioutil.ReadFile(CONFIG_FILE_PATH)
    if err != nil {
        return err
    }

    err = json.Unmarshal(file, &repo.Config)
    if err != nil {
        return err
    }
    return nil
}

func (repo Repository) SaveConfigToFile() error {
    jsonData, err := json.Marshal(repo.Config)
    if err != nil {
        return err
    }

    err = ioutil.WriteFile(CONFIG_FILE_PATH, jsonData, 0644)
    if err != nil {
        return err
    }

    return nil
}

func (repo *Repository) FindRelevantCommits() error {
    headRef, err := repo.GitRepo.LookupReference("HEAD")
    if err != nil {
        return err
    }

    headCommit, err := repo.GitRepo.LookupCommit(headRef.Target())
    if err != nil {
        return err
    }
    repo.HeadCommit = headCommit

    lastPushRef, err := repo.GitRepo.LookupReference(REF_S3_PUSH)
    if err != nil {
        return nil
    }

    lastPushCommit, err := repo.GitRepo.LookupCommit(lastPushRef.Target())
    if err != nil {
        return nil
    }
    repo.LastPushCommit = lastPushCommit

    return nil
}

func (repo *Repository) ReadGitModifiedFiles(scanner *bufio.Scanner, stop chan bool)  {
    for scanner.Scan() {
        repo.UnpushedFiles.Add(scanner.Text())
    }

    stop <- true
}

func (repo *Repository) FindCommitModifiedFiles(commit *gogit.Commit) error {
    cmd := exec.Command("git", "diff-tree", "--no-commit-id", "--name-only", "--root", commit.Id().String())
    out, err := cmd.StdoutPipe()
    if err != nil {
        return err
    }

    err = cmd.Start()
    if err != nil {
        return err
    }

    scanner := bufio.NewScanner(out)

    stop := make(chan bool)
    go repo.ReadGitModifiedFiles(scanner, stop)
    <-stop
    cmd.Wait()

    return nil
}

func (repo *Repository) FindUnpushedModifiedFiles() error {
    queue := []*gogit.Commit{};
    visited := mapset.NewSet();

    currentCommit := repo.HeadCommit;
    for currentCommit != nil && currentCommit.ParentCount() > 0 {
        if repo.LastPushCommit != nil && repo.LastPushCommit.Id().Equal(currentCommit.Id()) {
            break;
        }

        err := repo.FindCommitModifiedFiles(currentCommit)
        if err != nil {
            return err
        }

        for i := 0; i < currentCommit.ParentCount(); i++ {
            parentCommit := currentCommit.Parent(i)
            if !visited.Contains(parentCommit) {
                queue = append(queue, parentCommit)
            }
        }

        if len(queue) < 1 {
            break;
        }

        currentCommit = queue[0]
        queue = queue[1:]
    }

    return nil
}

func (repo Repository) UpdateGitLastPushRef() error {
    newLastPushRef := repo.HeadCommit.Id().String()
    cmd := exec.Command("git", "update-ref", REF_S3_PUSH, newLastPushRef)

    err := cmd.Start()
    if err != nil {
        return err
    }

    cmd.Wait()
    return nil
}

type S3Uploader struct {
    BucketName      *string
    S3Uploader      *s3manager.Uploader
}

func InitS3Uploader(config RepoConfig) *S3Uploader {
    uploader := new(S3Uploader)
    uploader.BucketName = aws.String(config.S3Bucket)

    s3config := aws.Config{Region: aws.String(config.S3Region)}
    s3uploader := s3manager.NewUploader(session.New(&s3config))
    uploader.S3Uploader = s3uploader

    return uploader
}

func (uploader S3Uploader) UploadFile(path string) error {
    file, err := os.Open(path)
    if err != nil {
        return err
    }

    result, err := uploader.S3Uploader.Upload(&s3manager.UploadInput{
        Body: file,
        Bucket: uploader.BucketName,
        Key: aws.String(path),
    })

    if err != nil {
        return err
    }

    fmt.Println(result.Location)
    return nil
}

func main() {
    repo, err := OpenRepository()
    if err != nil {
        fmt.Println(err)
        os.Exit(1)
    }
    repo.ReadConfigFile()

    flag.StringVar(&repo.Config.S3Bucket, "b", repo.Config.S3Bucket, "Destination S3 bucket name")
    flag.StringVar(&repo.Config.S3Region, "r", repo.Config.S3Region, "AWS region of destination bucket")
    saveConfig := flag.Bool("save", false, "Save destination region/bucket to config file")
    flag.Parse()

    if repo.Config.S3Bucket == "" {
        flag.Usage()
        os.Exit(1)
    } else if (repo.Config.S3Region == "") {
        flag.Usage()
        os.Exit(1)
    } else if (*saveConfig) {
        err := repo.SaveConfigToFile()
        if err != nil {
            fmt.Println("WARNING: Failed to save config to file: ", err)
        }
    }

    if err := repo.FindRelevantCommits(); err != nil {
        fmt.Println(err)
        os.Exit(1)
    }

    repo.FindUnpushedModifiedFiles();

    if repo.UnpushedFiles.Cardinality() == 0 {
        fmt.Println("No modified files to push")
        os.Exit(0)
    }

    uploader := InitS3Uploader(repo.Config)

    for filePath := range repo.UnpushedFiles.Iter() {
        fmt.Println("Uploading: " + filePath.(string))
        err := uploader.UploadFile(filePath.(string))
        if err != nil {
            fmt.Println(err)
            os.Exit(1)
        }
    }

    err = repo.UpdateGitLastPushRef()
    if err != nil {
        fmt.Println("Failed to update LAST_S3_PUSH ref with git: ", err)
    }
}<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"

	"code.google.com/p/go.crypto/ssh"
	"github.com/flynn/go-shlex"
)

const PrereceiveHookTmpl = `#!/bin/bash
set -eo pipefail; while read oldrev newrev refname; do
[[ $refname = "refs/heads/master" ]] && git archive $newrev | {{RECEIVER}} "$RECEIVE_REPO" "$newrev" | sed -$([[ $(uname) == "Darwin" ]] && echo l || echo u) "s/^/"$'\e[1G'"/"
done
`

var prereceiveHook []byte

var port = flag.String("p", "22", "port to listen on")
var repoPath = flag.String("r", "/tmp/repos", "path to repo cache")
var noAuth = flag.Bool("n", false, "disable client authentication")
var keys = flag.String("k", "", "pem file containing private keys (read from SSH_PRIVATE_KEYS by default)")

var authChecker []string

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %v [options] <authchecker> <receiver>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if flag.NArg() < 2 {
		flag.Usage()
		os.Exit(64)
	}
	var err error
	authChecker, err = shlex.Split(flag.Arg(0))
	if err != nil {
		log.Fatalln("Unable to parse authchecker command:", err)
	}

	receiver, err := shlex.Split(flag.Arg(1))
	if err != nil {
		log.Fatalln("Unable to parse receiver command:", err)
	}
	receiver[0], err = filepath.Abs(receiver[0])
	if err != nil {
		log.Fatalln("Invalid receiver path:", err)
	}
	prereceiveHook = []byte(strings.Replace(PrereceiveHookTmpl, "{{RECEIVER}}", strings.Join(receiver, " "), 1))

	var config *ssh.ServerConfig
	if *noAuth {
		config = &ssh.ServerConfig{NoClientAuth: true}
	} else {
		config = &ssh.ServerConfig{PublicKeyCallback: checkAuth}
	}

	if keyEnv := os.Getenv("SSH_PRIVATE_KEYS"); keyEnv != "" {
		if err := parseKeys(config, []byte(keyEnv)); err != nil {
			log.Fatalln("Failed to load private keys:", err)
		}
	} else {
		pemBytes, err := ioutil.ReadFile(*keys)
		if err != nil {
			log.Fatalln("Failed to load private keys:", err)
		}
		if err := parseKeys(config, pemBytes); err != nil {
			log.Fatalln("Failed to parse private keys:", err)
		}
	}

	if p := os.Getenv("PORT"); p != "" && *port == "22" {
		*port = p
	}
	listener, err := net.Listen("tcp", ":"+*port)
	if err != nil {
		log.Fatalln("Failed to listen for connections:", err)
	}
	for {
		// SSH connections just house multiplexed connections
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Failed to accept incoming connection:", err)
			continue
		}
		go handleConn(conn, config)
	}
}

func parseKeys(conf *ssh.ServerConfig, pemData []byte) error {
	var found bool
	for {
		var block *pem.Block
		block, pemData = pem.Decode(pemData)
		if block == nil {
			if !found {
				return errors.New("no private keys found")
			}
			return nil
		}
		if err := addKey(conf, block); err != nil {
			return err
		}
		found = true
	}
}

func addKey(conf *ssh.ServerConfig, block *pem.Block) (err error) {
	var key interface{}
	switch block.Type {
	case "RSA PRIVATE KEY":
		key, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	case "EC PRIVATE KEY":
		key, err = x509.ParseECPrivateKey(block.Bytes)
	case "DSA PRIVATE KEY":
		key, err = ssh.ParseDSAPrivateKey(block.Bytes)
	default:
		return fmt.Errorf("unsupported key type %q", block.Type)
	}
	if err != nil {
		return err
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return err
	}
	conf.AddHostKey(signer)
	return nil
}

func handleConn(conn net.Conn, conf *ssh.ServerConfig) {
	defer conn.Close()
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, conf)
	if err != nil {
		log.Println("Failed to handshake:", err)
		return
	}

	go ssh.DiscardRequests(reqs)

	for ch := range chans {
		if ch.ChannelType() != "session" {
			ch.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		go handleChannel(sshConn, ch)
	}
}

func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel) {
	ch, reqs, err := newChan.Accept()
	if err != nil {
		log.Println("newChan.Accept failed:", err)
		return
	}
	defer ch.Close()
	for req := range reqs {
		switch req.Type {
		case "exec":
			fail := func(at string, err error) {
				log.Printf("%s failed: %s", at, err)
				ch.Stderr().Write([]byte("Internal error.\n"))
			}
			if req.WantReply {
				req.Reply(true, nil)
			}
			cmdline := string(req.Payload[4:])
			cmdargs, err := shlex.Split(cmdline)
			if err != nil || len(cmdargs) != 2 {
				ch.Stderr().Write([]byte("Invalid arguments.\n"))
				return
			}
			if cmdargs[0] != "git-receive-pack" {
				ch.Stderr().Write([]byte("Only `git push` is supported.\n"))
				return
			}
			cmdargs[1] = strings.TrimSuffix(strings.TrimPrefix(cmdargs[1], "/"), ".git")
			if strings.Contains(cmdargs[1], "/") {
				ch.Stderr().Write([]byte("Invalid repo.\n"))
				return
			}

			if err := ensureCacheRepo(cmdargs[1]); err != nil {
				fail("ensureCacheRepo", err)
				return
			}
			cmd := exec.Command("git-shell", "-c", cmdargs[0]+" '"+cmdargs[1]+"'")
			cmd.Dir = *repoPath
			cmd.Env = append(os.Environ(),
				"RECEIVE_USER="+conn.User(),
				"RECEIVE_REPO="+cmdargs[1],
			)
			done, err := attachCmd(cmd, ch, ch.Stderr(), ch)
			if err != nil {
				fail("attachCmd", err)
				return
			}
			if err := cmd.Start(); err != nil {
				fail("cmd.Start", err)
				return
			}
			done.Wait()
			status, err := exitStatus(cmd.Wait())
			if err != nil {
				fail("exitStatus", err)
				return
			}
			if _, err := ch.SendRequest("exit-status", false, ssh.Marshal(&status)); err != nil {
				fail("sendExit", err)
			}
			return
		case "env":
			if req.WantReply {
				req.Reply(true, nil)
			}
		}
	}
}

var ErrUnauthorized = errors.New("gitreceive: user is unauthorized")

func checkAuth(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	status, err := exitStatus(exec.Command(authChecker[0],
		append(authChecker[1:], conn.User(), string(bytes.TrimSpace(ssh.MarshalAuthorizedKey(key))))...).Run())
	if err != nil {
		return nil, err
	}
	if status.Status == 0 {
		return nil, nil
	}
	return nil, ErrUnauthorized
}

func attachCmd(cmd *exec.Cmd, stdout, stderr io.Writer, stdin io.Reader) (*sync.WaitGroup, error) {
	var wg sync.WaitGroup
	wg.Add(2)

	stdinIn, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdoutOut, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderrOut, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}

	go func() {
		io.Copy(stdinIn, stdin)
		stdinIn.Close()
	}()
	go func() {
		io.Copy(stdout, stdoutOut)
		wg.Done()
	}()
	go func() {
		io.Copy(stderr, stderrOut)
		wg.Done()
	}()

	return &wg, nil
}

type exitStatusMsg struct {
	Status uint32
}

func exitStatus(err error) (exitStatusMsg, error) {
	if err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			// There is no platform independent way to retrieve
			// the exit code, but the following will work on Unix
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				return exitStatusMsg{uint32(status.ExitStatus())}, nil
			}
		}
		return exitStatusMsg{0}, err
	}
	return exitStatusMsg{0}, nil
}

var cacheMtx sync.Mutex

func ensureCacheRepo(path string) error {
	cacheMtx.Lock()
	defer cacheMtx.Unlock()

	cachePath := *repoPath + "/" + path
	if _, err := os.Stat(cachePath); os.IsNotExist(err) {
		os.MkdirAll(cachePath, 0755)
		cmd := exec.Command("git", "init", "--bare")
		cmd.Dir = cachePath
		err = cmd.Run()
		if err != nil {
			return err
		}
	}
	return ioutil.WriteFile(cachePath+"/hooks/pre-receive", prereceiveHook, 0755)
}
<commit_msg>gitreceived: Merge pull request #14 from miracle2k/master<commit_after>package main

import (
	"bytes"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"

	"code.google.com/p/go.crypto/ssh"
	"github.com/flynn/go-shlex"
)

const PrereceiveHookTmpl = `#!/bin/bash
set -eo pipefail; while read oldrev newrev refname; do
[[ $refname = "refs/heads/master" ]] && git archive $newrev | {{RECEIVER}} "$RECEIVE_REPO" "$newrev" | sed -$([[ $(uname) == "Darwin" ]] && echo l || echo u) "s/^/"$'\e[1G'"/"
done
`

var prereceiveHook []byte

var port = flag.String("p", "22", "port to listen on")
var repoPath = flag.String("r", "/tmp/repos", "path to repo cache")
var noAuth = flag.Bool("n", false, "disable client authentication")
var keys = flag.String("k", "", "pem file containing private keys (read from SSH_PRIVATE_KEYS by default)")

var authChecker []string

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %v [options] <authchecker> <receiver>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if flag.NArg() < 2 {
		flag.Usage()
		os.Exit(64)
	}
	var err error
	authChecker, err = shlex.Split(flag.Arg(0))
	if err != nil {
		log.Fatalln("Unable to parse authchecker command:", err)
	}

	receiver, err := shlex.Split(flag.Arg(1))
	if err != nil {
		log.Fatalln("Unable to parse receiver command:", err)
	}
	receiver[0], err = filepath.Abs(receiver[0])
	if err != nil {
		log.Fatalln("Invalid receiver path:", err)
	}
	prereceiveHook = []byte(strings.Replace(PrereceiveHookTmpl, "{{RECEIVER}}", strings.Join(receiver, " "), 1))

	var config *ssh.ServerConfig
	if *noAuth {
		config = &ssh.ServerConfig{NoClientAuth: true}
	} else {
		config = &ssh.ServerConfig{PublicKeyCallback: checkAuth}
	}

	if keyEnv := os.Getenv("SSH_PRIVATE_KEYS"); keyEnv != "" {
		if err := parseKeys(config, []byte(keyEnv)); err != nil {
			log.Fatalln("Failed to load private keys:", err)
		}
	} else {
		pemBytes, err := ioutil.ReadFile(*keys)
		if err != nil {
			log.Fatalln("Failed to load private keys:", err)
		}
		if err := parseKeys(config, pemBytes); err != nil {
			log.Fatalln("Failed to parse private keys:", err)
		}
	}

	if p := os.Getenv("PORT"); p != "" && *port == "22" {
		*port = p
	}
	listener, err := net.Listen("tcp", ":"+*port)
	if err != nil {
		log.Fatalln("Failed to listen for connections:", err)
	}
	for {
		// SSH connections just house multiplexed connections
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Failed to accept incoming connection:", err)
			continue
		}
		go handleConn(conn, config)
	}
}

func parseKeys(conf *ssh.ServerConfig, pemData []byte) error {
	var found bool
	for {
		var block *pem.Block
		block, pemData = pem.Decode(pemData)
		if block == nil {
			if !found {
				return errors.New("no private keys found")
			}
			return nil
		}
		if err := addKey(conf, block); err != nil {
			return err
		}
		found = true
	}
}

func addKey(conf *ssh.ServerConfig, block *pem.Block) (err error) {
	var key interface{}
	switch block.Type {
	case "RSA PRIVATE KEY":
		key, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	case "EC PRIVATE KEY":
		key, err = x509.ParseECPrivateKey(block.Bytes)
	case "DSA PRIVATE KEY":
		key, err = ssh.ParseDSAPrivateKey(block.Bytes)
	default:
		return fmt.Errorf("unsupported key type %q", block.Type)
	}
	if err != nil {
		return err
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return err
	}
	conf.AddHostKey(signer)
	return nil
}

func handleConn(conn net.Conn, conf *ssh.ServerConfig) {
	defer conn.Close()
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, conf)
	if err != nil {
		log.Println("Failed to handshake:", err)
		return
	}

	go ssh.DiscardRequests(reqs)

	for ch := range chans {
		if ch.ChannelType() != "session" {
			ch.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		go handleChannel(sshConn, ch)
	}
}

func handleChannel(conn *ssh.ServerConn, newChan ssh.NewChannel) {
	ch, reqs, err := newChan.Accept()
	if err != nil {
		log.Println("newChan.Accept failed:", err)
		return
	}
	defer ch.Close()
	for req := range reqs {
		switch req.Type {
		case "exec":
			fail := func(at string, err error) {
				log.Printf("%s failed: %s", at, err)
				ch.Stderr().Write([]byte("Internal error.\n"))
			}
			if req.WantReply {
				req.Reply(true, nil)
			}
			cmdline := string(req.Payload[4:])
			cmdargs, err := shlex.Split(cmdline)
			if err != nil || len(cmdargs) != 2 {
				ch.Stderr().Write([]byte("Invalid arguments.\n"))
				return
			}
			if cmdargs[0] != "git-receive-pack" {
				ch.Stderr().Write([]byte("Only `git push` is supported.\n"))
				return
			}
			cmdargs[1] = strings.TrimSuffix(strings.TrimPrefix(cmdargs[1], "/"), ".git")
			if strings.Contains(cmdargs[1], "..") {
				ch.Stderr().Write([]byte("Invalid repo.\n"))
				return
			}

			if err := ensureCacheRepo(cmdargs[1]); err != nil {
				fail("ensureCacheRepo", err)
				return
			}
			cmd := exec.Command("git-shell", "-c", cmdargs[0]+" '"+cmdargs[1]+"'")
			cmd.Dir = *repoPath
			cmd.Env = append(os.Environ(),
				"RECEIVE_USER="+conn.User(),
				"RECEIVE_REPO="+cmdargs[1],
			)
			done, err := attachCmd(cmd, ch, ch.Stderr(), ch)
			if err != nil {
				fail("attachCmd", err)
				return
			}
			if err := cmd.Start(); err != nil {
				fail("cmd.Start", err)
				return
			}
			done.Wait()
			status, err := exitStatus(cmd.Wait())
			if err != nil {
				fail("exitStatus", err)
				return
			}
			if _, err := ch.SendRequest("exit-status", false, ssh.Marshal(&status)); err != nil {
				fail("sendExit", err)
			}
			return
		case "env":
			if req.WantReply {
				req.Reply(true, nil)
			}
		}
	}
}

var ErrUnauthorized = errors.New("gitreceive: user is unauthorized")

func checkAuth(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	status, err := exitStatus(exec.Command(authChecker[0],
		append(authChecker[1:], conn.User(), string(bytes.TrimSpace(ssh.MarshalAuthorizedKey(key))))...).Run())
	if err != nil {
		return nil, err
	}
	if status.Status == 0 {
		return nil, nil
	}
	return nil, ErrUnauthorized
}

func attachCmd(cmd *exec.Cmd, stdout, stderr io.Writer, stdin io.Reader) (*sync.WaitGroup, error) {
	var wg sync.WaitGroup
	wg.Add(2)

	stdinIn, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdoutOut, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderrOut, err := cmd.StderrPipe()
	if err != nil {
		return nil, err
	}

	go func() {
		io.Copy(stdinIn, stdin)
		stdinIn.Close()
	}()
	go func() {
		io.Copy(stdout, stdoutOut)
		wg.Done()
	}()
	go func() {
		io.Copy(stderr, stderrOut)
		wg.Done()
	}()

	return &wg, nil
}

type exitStatusMsg struct {
	Status uint32
}

func exitStatus(err error) (exitStatusMsg, error) {
	if err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			// There is no platform independent way to retrieve
			// the exit code, but the following will work on Unix
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				return exitStatusMsg{uint32(status.ExitStatus())}, nil
			}
		}
		return exitStatusMsg{0}, err
	}
	return exitStatusMsg{0}, nil
}

var cacheMtx sync.Mutex

func ensureCacheRepo(path string) error {
	cacheMtx.Lock()
	defer cacheMtx.Unlock()

	cachePath := *repoPath + "/" + path
	if _, err := os.Stat(cachePath); os.IsNotExist(err) {
		os.MkdirAll(cachePath, 0755)
		cmd := exec.Command("git", "init", "--bare")
		cmd.Dir = cachePath
		err = cmd.Run()
		if err != nil {
			return err
		}
	}
	return ioutil.WriteFile(cachePath+"/hooks/pre-receive", prereceiveHook, 0755)
}
<|endoftext|>


<commit_before>// Copyright 2011 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 main

import (
	"fmt"
	"go/build"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
)

var cmdTool = &Command{
	Run:       runTool,
	UsageLine: "tool command [args...]",
	Short:     "run specified go tool",
	Long: `
Tool runs the go tool command identified by the arguments.
With no arguments it prints the list of known tools.

For more about each tool command, see 'go tool command -h'.
`,
}

var (
	toolGOOS      = runtime.GOOS
	toolGOARCH    = runtime.GOARCH
	toolIsWindows = toolGOOS == "windows"
	toolDir       = filepath.Join(build.Path[0].Path, "bin", "go-tool")
)

const toolWindowsExtension = ".exe"

func tool(name string) string {
	p := filepath.Join(toolDir, name)
	if toolIsWindows {
		p += toolWindowsExtension
	}
	return p
}

func runTool(cmd *Command, args []string) {
	if len(args) == 0 {
		listTools()
		return
	}
	toolName := args[0]
	// The tool name must be lower-case letters and numbers.
	for _, c := range toolName {
		switch {
		case 'a' <= c && c <= 'z', '0' <= c && c <= '9':
		default:
			fmt.Fprintf(os.Stderr, "go tool: bad tool name %q\n", tool)
			setExitStatus(2)
			return
		}
	}
	toolPath := tool(toolName)
	// Give a nice message if there is no tool with that name.
	if _, err := os.Stat(toolPath); err != nil {
		fmt.Fprintf(os.Stderr, "go tool: no such tool %q\n", tool)
		setExitStatus(3)
		return
	}
	toolCmd := &exec.Cmd{
		Path:   toolPath,
		Args:   args,
		Stdout: os.Stdout,
		Stderr: os.Stderr,
	}
	err := toolCmd.Run()
	if err != nil {
		fmt.Fprintf(os.Stderr, "go tool %s: %s\n", toolName, err)
		setExitStatus(1)
		return
	}
}

// listTools prints a list of the available tools in the go-tools directory.
func listTools() {
	f, err := os.Open(toolDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "go tool: no tool directory: %s\n", err)
		setExitStatus(2)
		return
	}
	defer f.Close()
	names, err := f.Readdirnames(-1)
	if err != nil {
		fmt.Fprintf(os.Stderr, "go tool: can't read directory: %s\n", err)
		setExitStatus(2)
		return
	}
	sort.Strings(names)
	for _, name := range names {
		// Unify presentation by going to lower case.
		name = strings.ToLower(name)
		// If it's windows, don't show the .exe suffix.
		if toolIsWindows && strings.HasSuffix(name, toolWindowsExtension) {
			name = name[:len(name)-len(toolWindowsExtension)]
		}
		fmt.Println(name)
	}
}
<commit_msg>cmd/go: fix error message on non-existing tools.<commit_after>// Copyright 2011 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 main

import (
	"fmt"
	"go/build"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
)

var cmdTool = &Command{
	Run:       runTool,
	UsageLine: "tool command [args...]",
	Short:     "run specified go tool",
	Long: `
Tool runs the go tool command identified by the arguments.
With no arguments it prints the list of known tools.

For more about each tool command, see 'go tool command -h'.
`,
}

var (
	toolGOOS      = runtime.GOOS
	toolGOARCH    = runtime.GOARCH
	toolIsWindows = toolGOOS == "windows"
	toolDir       = filepath.Join(build.Path[0].Path, "bin", "go-tool")
)

const toolWindowsExtension = ".exe"

func tool(name string) string {
	p := filepath.Join(toolDir, name)
	if toolIsWindows {
		p += toolWindowsExtension
	}
	return p
}

func runTool(cmd *Command, args []string) {
	if len(args) == 0 {
		listTools()
		return
	}
	toolName := args[0]
	// The tool name must be lower-case letters and numbers.
	for _, c := range toolName {
		switch {
		case 'a' <= c && c <= 'z', '0' <= c && c <= '9':
		default:
			fmt.Fprintf(os.Stderr, "go tool: bad tool name %q\n", toolName)
			setExitStatus(2)
			return
		}
	}
	toolPath := tool(toolName)
	// Give a nice message if there is no tool with that name.
	if _, err := os.Stat(toolPath); err != nil {
		fmt.Fprintf(os.Stderr, "go tool: no such tool %q\n", toolName)
		setExitStatus(3)
		return
	}
	toolCmd := &exec.Cmd{
		Path:   toolPath,
		Args:   args,
		Stdout: os.Stdout,
		Stderr: os.Stderr,
	}
	err := toolCmd.Run()
	if err != nil {
		fmt.Fprintf(os.Stderr, "go tool %s: %s\n", toolName, err)
		setExitStatus(1)
		return
	}
}

// listTools prints a list of the available tools in the go-tools directory.
func listTools() {
	f, err := os.Open(toolDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "go tool: no tool directory: %s\n", err)
		setExitStatus(2)
		return
	}
	defer f.Close()
	names, err := f.Readdirnames(-1)
	if err != nil {
		fmt.Fprintf(os.Stderr, "go tool: can't read directory: %s\n", err)
		setExitStatus(2)
		return
	}
	sort.Strings(names)
	for _, name := range names {
		// Unify presentation by going to lower case.
		name = strings.ToLower(name)
		// If it's windows, don't show the .exe suffix.
		if toolIsWindows && strings.HasSuffix(name, toolWindowsExtension) {
			name = name[:len(name)-len(toolWindowsExtension)]
		}
		fmt.Println(name)
	}
}
<|endoftext|>


<commit_before>//
// Copyright 2019 Google Inc. All Rights Reserved.
//
// 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 client creates a tunnel client to proxy incoming connections
// over a grpc transport.
package client

import (
	"context"
	"fmt"
	"io"
	"log"
	"net"

	"github.com/openconfig/grpctunnel/bidi"
	"github.com/openconfig/grpctunnel/tunnel"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	tpb "github.com/openconfig/grpctunnel/proto/tunnel"
)

// Config defines the parameters to run a tunnel client.
type Config struct {
	TunnelAddress, DialAddress, ListenAddress, CertFile, Target, TargetType string
	Bridge                                                                  bool
}

func listen(ctx context.Context, c *tunnel.Client, listenAddress string, target tunnel.Target) error {
	l, err := net.Listen("tcp", listenAddress)
	if err != nil {
		return fmt.Errorf("failed to listen: %s: %v", listenAddress, err)
	}
	defer l.Close()

	errCh := make(chan error)
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				select {
				case errCh <- fmt.Errorf("failed to accept connection: %v", err):
				default:
				}
				return
			}
			// Errors from this goroutine will be logged only, because we don't want an
			// underlying stream issue to tear the server down
			go func(conn net.Conn) {
				defer conn.Close()
				session, err := c.NewSession(target)
				if err != nil {
					log.Printf("error from new session: %v", err)
					return
				}
				if err = bidi.Copy(session, conn); err != nil {
					log.Printf("error from bidi copy: %v", err)
				}
			}(conn)
		}
	}()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-errCh:
		return err
	}
}

// Run starts a tunnel client, connecting to the tunnel server via the provided tunnel address.
// The client uses the target to identify whether it can handle the target (u) sent by the server.
func Run(ctx context.Context, conf Config) error {
	opts := []grpc.DialOption{grpc.WithDefaultCallOptions()}
	if conf.CertFile == "" {
		opts = append(opts, grpc.WithInsecure())
	} else {
		creds, err := credentials.NewClientTLSFromFile(conf.CertFile, "")
		if err != nil {
			return fmt.Errorf("failed to load credentials: %v", err)
		}
		opts = append(opts, grpc.WithTransportCredentials(creds))
	}
	clientConn, err := grpc.Dial(conf.TunnelAddress, opts...)
	if err != nil {
		return fmt.Errorf("grpc dial error: %v", err)
	}
	defer clientConn.Close()

	registerHandler := func(t tunnel.Target) error {
		if t.ID != conf.Target {
			return fmt.Errorf("client cannot handle: %s", t.ID)
		}
		return nil
	}

	handler := func(_ tunnel.Target, i io.ReadWriteCloser) error {
		conn, err := net.Dial("tcp", conf.DialAddress)
		if err != nil {
			log.Printf("Error dialing client: %v", err)
			return err
		}

		if err = bidi.Copy(i, conn); err != nil {
			// Logging this error only as we don't want the client to stop because an
			// underlying stream had an issue
			log.Printf("Copy error: %v", err)
		}

		return nil
	}

	peerAddCh := make(chan tunnel.Target, 1)
	peerAddHandler := func(t tunnel.Target) error {
		peerAddCh <- t
		return nil
	}

	peerDelCh := make(chan tunnel.Target, 1)
	peerDelHandler := func(t tunnel.Target) error {
		peerDelCh <- t
		return nil
	}

	targets := make(map[tunnel.Target]struct{})
	t := tunnel.Target{ID: conf.Target, Type: conf.TargetType}
	targets[t] = struct{}{}
	client, err := tunnel.NewClient(tpb.NewTunnelClient(clientConn), tunnel.ClientConfig{
		RegisterHandler: registerHandler,
		Handler:         handler,
		Subscriptions:   []string{conf.TargetType},
		PeerAddHandler:  peerAddHandler,
		PeerDelHandler:  peerDelHandler,
	}, targets)

	if err != nil {
		return fmt.Errorf("failed to create tunnel client: %v", err)
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	errCh := make(chan error, 2)
	go func() {
		err := client.Run(ctx)
		if err != nil {
			errCh <- err
		}
	}()

	// If running bridge mode, pick peer targets sequentially.
	if conf.Bridge {
		cancels := make(map[tunnel.Target]func())
		select {
		case target := <-peerAddCh:
			ctx, cancels[target] = context.WithCancel(ctx)
			if err := listen(ctx, client, conf.ListenAddress, target); err != nil {
				errCh <- err
				cancels[target]()
			}
		case target := <-peerDelCh:
			cancels[target]()
			delete(cancels, target)
		}
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-errCh:
		return err
	}

}
<commit_msg>clean up bridge flag<commit_after>//
// Copyright 2019 Google Inc. All Rights Reserved.
//
// 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 client creates a tunnel client to proxy incoming connections
// over a grpc transport.
package client

import (
	"context"
	"fmt"
	"io"
	"log"
	"net"

	"github.com/openconfig/grpctunnel/bidi"
	"github.com/openconfig/grpctunnel/tunnel"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	tpb "github.com/openconfig/grpctunnel/proto/tunnel"
)

// Config defines the parameters to run a tunnel client.
type Config struct {
	TunnelAddress, DialAddress, ListenAddress, CertFile, Target, TargetType string
}

func listen(ctx context.Context, c *tunnel.Client, listenAddress string, target tunnel.Target) error {
	l, err := net.Listen("tcp", listenAddress)
	if err != nil {
		return fmt.Errorf("failed to listen: %s: %v", listenAddress, err)
	}
	defer l.Close()

	errCh := make(chan error)
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				select {
				case errCh <- fmt.Errorf("failed to accept connection: %v", err):
				default:
				}
				return
			}
			// Errors from this goroutine will be logged only, because we don't want an
			// underlying stream issue to tear the server down
			go func(conn net.Conn) {
				defer conn.Close()
				session, err := c.NewSession(target)
				if err != nil {
					log.Printf("error from new session: %v", err)
					return
				}
				if err = bidi.Copy(session, conn); err != nil {
					log.Printf("error from bidi copy: %v", err)
				}
			}(conn)
		}
	}()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-errCh:
		return err
	}
}

// Run starts a tunnel client, connecting to the tunnel server via the provided tunnel address.
// The client uses the target to identify whether it can handle the target (u) sent by the server.
func Run(ctx context.Context, conf Config) error {
	opts := []grpc.DialOption{grpc.WithDefaultCallOptions()}
	if conf.CertFile == "" {
		opts = append(opts, grpc.WithInsecure())
	} else {
		creds, err := credentials.NewClientTLSFromFile(conf.CertFile, "")
		if err != nil {
			return fmt.Errorf("failed to load credentials: %v", err)
		}
		opts = append(opts, grpc.WithTransportCredentials(creds))
	}
	clientConn, err := grpc.Dial(conf.TunnelAddress, opts...)
	if err != nil {
		return fmt.Errorf("grpc dial error: %v", err)
	}
	defer clientConn.Close()

	registerHandler := func(t tunnel.Target) error {
		if t.ID != conf.Target {
			return fmt.Errorf("client cannot handle: %s", t.ID)
		}
		return nil
	}

	handler := func(_ tunnel.Target, i io.ReadWriteCloser) error {
		conn, err := net.Dial("tcp", conf.DialAddress)
		if err != nil {
			log.Printf("Error dialing client: %v", err)
			return err
		}

		if err = bidi.Copy(i, conn); err != nil {
			// Logging this error only as we don't want the client to stop because an
			// underlying stream had an issue
			log.Printf("Copy error: %v", err)
		}

		return nil
	}

	peerAddCh := make(chan tunnel.Target, 1)
	peerAddHandler := func(t tunnel.Target) error {
		peerAddCh <- t
		return nil
	}

	peerDelCh := make(chan tunnel.Target, 1)
	peerDelHandler := func(t tunnel.Target) error {
		peerDelCh <- t
		return nil
	}

	targets := make(map[tunnel.Target]struct{})
	t := tunnel.Target{ID: conf.Target, Type: conf.TargetType}
	targets[t] = struct{}{}
	client, err := tunnel.NewClient(tpb.NewTunnelClient(clientConn), tunnel.ClientConfig{
		RegisterHandler: registerHandler,
		Handler:         handler,
		Subscriptions:   []string{conf.TargetType},
		PeerAddHandler:  peerAddHandler,
		PeerDelHandler:  peerDelHandler,
	}, targets)

	if err != nil {
		return fmt.Errorf("failed to create tunnel client: %v", err)
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	errCh := make(chan error, 2)
	go func() {
		err := client.Run(ctx)
		if err != nil {
			errCh <- err
		}
	}()

	// listen for any request to create a new session
	cancels := make(map[tunnel.Target]func())
	select {
	case target := <-peerAddCh:
		ctx, cancels[target] = context.WithCancel(ctx)
		if err := listen(ctx, client, conf.ListenAddress, target); err != nil {
			errCh <- err
			cancels[target]()
		}
	case target := <-peerDelCh:
		cancels[target]()
		delete(cancels, target)
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-errCh:
		return err
	}

}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"errors"
	"fmt"

	"github.com/1and1/soma/internal/stmt"
	"github.com/1and1/soma/lib/proto"
	log "github.com/Sirupsen/logrus"
	uuid "github.com/satori/go.uuid"
)

type somaNodeRequest struct {
	action string
	user   string
	Node   proto.Node
	reply  chan somaResult
}

type somaNodeResult struct {
	ResultError error
	Node        proto.Node
}

func (a *somaNodeResult) SomaAppendError(r *somaResult, err error) {
	if err != nil {
		r.Nodes = append(r.Nodes, somaNodeResult{ResultError: err})
	}
}

func (a *somaNodeResult) SomaAppendResult(r *somaResult) {
	r.Nodes = append(r.Nodes, *a)
}

/* Read Access
 *
 */
type somaNodeReadHandler struct {
	input     chan somaNodeRequest
	shutdown  chan bool
	conn      *sql.DB
	list_stmt *sql.Stmt
	show_stmt *sql.Stmt
	conf_stmt *sql.Stmt
	sync_stmt *sql.Stmt
	ponc_stmt *sql.Stmt
	psvc_stmt *sql.Stmt
	psys_stmt *sql.Stmt
	pcst_stmt *sql.Stmt
	appLog    *log.Logger
	reqLog    *log.Logger
	errLog    *log.Logger
}

func (r *somaNodeReadHandler) run() {
	var err error

	for statement, prepStmt := range map[string]*sql.Stmt{
		stmt.ListNodes:       r.list_stmt,
		stmt.ShowNodes:       r.show_stmt,
		stmt.ShowConfigNodes: r.conf_stmt,
		stmt.SyncNodes:       r.sync_stmt,
		stmt.NodeOncProps:    r.ponc_stmt,
		stmt.NodeSvcProps:    r.psvc_stmt,
		stmt.NodeSysProps:    r.psys_stmt,
		stmt.NodeCstProps:    r.pcst_stmt,
	} {
		if prepStmt, err = r.conn.Prepare(statement); err != nil {
			r.errLog.Fatal(`node`, err, statement)
		}
		defer prepStmt.Close()
	}

runloop:
	for {
		select {
		case <-r.shutdown:
			break runloop
		case req := <-r.input:
			go func() {
				r.process(&req)
			}()
		}
	}
}

func (r *somaNodeReadHandler) process(q *somaNodeRequest) {
	var (
		nodeId, nodeName, nodeTeam, nodeServer, nodeState    string
		bucketId, repositoryId, instanceId, sourceInstanceId string
		view, oncallId, oncallName, serviceName, customId    string
		systemProp, value, customProp                        string
		nodeAsset                                            int
		nodeOnline, nodeDeleted                              bool
		rows                                                 *sql.Rows
		err                                                  error
	)
	result := somaResult{}

	switch q.action {
	case "list":
		r.reqLog.Printf("R: node/list")
		rows, err = r.list_stmt.Query()
		if result.SetRequestError(err) {
			goto dispatch
		}
		defer rows.Close()

		for rows.Next() {
			err = rows.Scan(&nodeId, &nodeName)
			result.Append(err, &somaNodeResult{
				Node: proto.Node{
					Id:   nodeId,
					Name: nodeName,
				},
			})
		}
	case `sync`:
		r.reqLog.Printf(`R: node/sync`)
		rows, err = r.sync_stmt.Query()
		if result.SetRequestError(err) {
			goto dispatch
		}
		defer rows.Close()

		for rows.Next() {
			err = rows.Scan(
				&nodeId,
				&nodeAsset,
				&nodeName,
				&nodeTeam,
				&nodeServer,
				&nodeOnline,
				&nodeDeleted,
			)
			result.Append(err, &somaNodeResult{
				Node: proto.Node{
					Id:        nodeId,
					AssetId:   uint64(nodeAsset),
					Name:      nodeName,
					TeamId:    nodeTeam,
					ServerId:  nodeServer,
					IsOnline:  nodeOnline,
					IsDeleted: nodeDeleted,
				},
			})
		}
	case "show":
		r.reqLog.Printf("R: node/show")
		err = r.show_stmt.QueryRow(q.Node.Id).Scan(
			&nodeId,
			&nodeAsset,
			&nodeName,
			&nodeTeam,
			&nodeServer,
			&nodeState,
			&nodeOnline,
			&nodeDeleted,
		)
		if err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			goto dispatch
		}
		node := proto.Node{
			Id:        nodeId,
			AssetId:   uint64(nodeAsset),
			Name:      nodeName,
			TeamId:    nodeTeam,
			ServerId:  nodeServer,
			State:     nodeState,
			IsOnline:  nodeOnline,
			IsDeleted: nodeDeleted,
		}

		// add configuration data
		err = r.conf_stmt.QueryRow(q.Node.Id).Scan(
			&nodeId,
			&nodeName,
			&bucketId,
			&repositoryId,
		)
		if err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			goto dispatch
		}
		node.Config = &proto.NodeConfig{
			RepositoryId: repositoryId,
			BucketId:     bucketId,
		}
		node.Properties = &[]proto.Property{}

		// oncall properties
		rows, err = r.ponc_stmt.Query(q.Node.Id)
		if result.SetRequestError(err) {
			goto dispatch
		}
		for rows.Next() {
			if err = rows.Scan(
				&instanceId,
				&sourceInstanceId,
				&view,
				&oncallId,
				&oncallName,
			); result.SetRequestError(err) {
				rows.Close()
				goto dispatch
			}
			*node.Properties = append(
				*node.Properties,
				proto.Property{
					Type:             `oncall`,
					RepositoryId:     repositoryId,
					BucketId:         bucketId,
					InstanceId:       instanceId,
					SourceInstanceId: sourceInstanceId,
					View:             view,
					Oncall: &proto.PropertyOncall{
						Id:   oncallId,
						Name: oncallName,
					},
				},
			)
		}
		if err = rows.Err(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}

		// service properties
		rows, err = r.psvc_stmt.Query(q.Node.Id)
		for rows.Next() {
			if err = rows.Scan(
				&instanceId,
				&sourceInstanceId,
				&view,
				&serviceName,
			); result.SetRequestError(err) {
				rows.Close()
				goto dispatch
			}
			*node.Properties = append(
				*node.Properties,
				proto.Property{
					Type:             `service`,
					RepositoryId:     repositoryId,
					BucketId:         bucketId,
					InstanceId:       instanceId,
					SourceInstanceId: sourceInstanceId,
					View:             view,
					Service: &proto.PropertyService{
						Name: serviceName,
					},
				},
			)
		}
		if err = rows.Err(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}

		// system properties
		rows, err = r.psys_stmt.Query(q.Node.Id)
		for rows.Next() {
			if err = rows.Scan(
				&instanceId,
				&sourceInstanceId,
				&view,
				&systemProp,
				&value,
			); result.SetRequestError(err) {
				rows.Close()
				goto dispatch
			}
			*node.Properties = append(
				*node.Properties,
				proto.Property{
					Type:             `system`,
					RepositoryId:     repositoryId,
					BucketId:         bucketId,
					InstanceId:       instanceId,
					SourceInstanceId: sourceInstanceId,
					View:             view,
					System: &proto.PropertySystem{
						Name:  systemProp,
						Value: value,
					},
				},
			)
		}
		if err = rows.Err(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}

		// custom properties
		rows, err = r.pcst_stmt.Query(q.Node.Id)
		for rows.Next() {
			if err = rows.Scan(
				&instanceId,
				&sourceInstanceId,
				&view,
				&customId,
				&value,
				&customProp,
			); result.SetRequestError(err) {
				rows.Close()
				goto dispatch
			}
			*node.Properties = append(
				*node.Properties,
				proto.Property{
					Type:             `custom`,
					RepositoryId:     repositoryId,
					BucketId:         bucketId,
					InstanceId:       instanceId,
					SourceInstanceId: sourceInstanceId,
					View:             view,
					Custom: &proto.PropertyCustom{
						Id:    customId,
						Name:  customProp,
						Value: value,
					},
				},
			)
		}
		if err = rows.Err(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}

		result.Append(err, &somaNodeResult{
			Node: node,
		})
	case "get_config":
		r.reqLog.Printf("R: node/get_config")
		err = r.conf_stmt.QueryRow(q.Node.Id).Scan(
			&nodeId,
			&nodeName,
			&bucketId,
			&repositoryId,
		)
		if err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			goto dispatch
		}

		result.Append(err, &somaNodeResult{
			Node: proto.Node{
				Id:   nodeId,
				Name: nodeName,
				Config: &proto.NodeConfig{
					RepositoryId: repositoryId,
					BucketId:     bucketId,
				},
			},
		})
	default:
		result.SetNotImplemented()
	}
dispatch:
	q.reply <- result
}

/* Write Access
 *
 */
type somaNodeWriteHandler struct {
	input    chan somaNodeRequest
	shutdown chan bool
	conn     *sql.DB
	add_stmt *sql.Stmt
	del_stmt *sql.Stmt
	prg_stmt *sql.Stmt
	upd_stmt *sql.Stmt
	appLog   *log.Logger
	reqLog   *log.Logger
	errLog   *log.Logger
}

func (w *somaNodeWriteHandler) run() {
	var err error

	for statement, prepStmt := range map[string]*sql.Stmt{
		stmt.NodeAdd:    w.add_stmt,
		stmt.NodeUpdate: w.upd_stmt,
		stmt.NodeDel:    w.del_stmt,
		stmt.NodePurge:  w.prg_stmt,
	} {
		if prepStmt, err = w.conn.Prepare(statement); err != nil {
			w.errLog.Fatal(`node`, err, statement)
		}
		defer prepStmt.Close()
	}

runloop:
	for {
		select {
		case <-w.shutdown:
			break runloop
		case req := <-w.input:
			go func() {
				w.process(&req)
			}()
		}
	}
}

func (w *somaNodeWriteHandler) process(q *somaNodeRequest) {
	var res sql.Result
	var err error
	result := somaResult{}

	switch q.action {
	case "add":
		w.reqLog.Printf("R: node/add for %s", q.Node.Name)
		id := uuid.NewV4()
		if q.Node.ServerId == "" {
			q.Node.ServerId = "00000000-0000-0000-0000-000000000000"
		}
		res, err = w.add_stmt.Exec(
			id.String(),
			q.Node.AssetId,
			q.Node.Name,
			q.Node.TeamId,
			q.Node.ServerId,
			q.Node.State,
			q.Node.IsOnline,
			false,
			q.user,
		)
		q.Node.Id = id.String()
	case `update`:
		w.reqLog.Printf("R: node/update for %s", q.Node.Id)
		res, err = w.upd_stmt.Exec(
			q.Node.AssetId,
			q.Node.Name,
			q.Node.TeamId,
			q.Node.ServerId,
			q.Node.IsOnline,
			q.Node.IsDeleted,
			q.Node.Id,
		)
		// TODO what has to be done for this undeployment?
	case "delete":
		w.reqLog.Printf("R: node/delete for %s", q.Node.Id)
		res, err = w.del_stmt.Exec(
			q.Node.Id,
		)
		// TODO trigger undeployment
	case "purge":
		w.reqLog.Printf("R: node/purge for %s", q.Node.Id)
		res, err = w.prg_stmt.Exec(
			q.Node.Id,
		)
	default:
		w.reqLog.Printf("R: unimplemented node/%s", q.action)
		result.SetNotImplemented()
		q.reply <- result
		return
	}
	if result.SetRequestError(err) {
		q.reply <- result
		return
	}

	rowCnt, _ := res.RowsAffected()
	switch {
	case rowCnt == 0:
		result.Append(errors.New("No rows affected"), &somaNodeResult{})
	case rowCnt > 1:
		result.Append(fmt.Errorf("Too many rows affected: %d", rowCnt),
			&somaNodeResult{})
	default:
		result.Append(nil, &somaNodeResult{
			Node: q.Node,
		})
	}
	q.reply <- result
}

/* Ops Access
 */
func (r *somaNodeReadHandler) shutdownNow() {
	r.shutdown <- true
}

func (w *somaNodeWriteHandler) shutdownNow() {
	w.shutdown <- true
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Add checkconfig/instance info to node/show<commit_after>package main

import (
	"database/sql"
	"errors"
	"fmt"

	"github.com/1and1/soma/internal/stmt"
	"github.com/1and1/soma/lib/proto"
	log "github.com/Sirupsen/logrus"
	uuid "github.com/satori/go.uuid"
)

type somaNodeRequest struct {
	action string
	user   string
	Node   proto.Node
	reply  chan somaResult
}

type somaNodeResult struct {
	ResultError error
	Node        proto.Node
}

func (a *somaNodeResult) SomaAppendError(r *somaResult, err error) {
	if err != nil {
		r.Nodes = append(r.Nodes, somaNodeResult{ResultError: err})
	}
}

func (a *somaNodeResult) SomaAppendResult(r *somaResult) {
	r.Nodes = append(r.Nodes, *a)
}

/* Read Access
 *
 */
type somaNodeReadHandler struct {
	input     chan somaNodeRequest
	shutdown  chan bool
	conn      *sql.DB
	list_stmt *sql.Stmt
	show_stmt *sql.Stmt
	conf_stmt *sql.Stmt
	sync_stmt *sql.Stmt
	ponc_stmt *sql.Stmt
	psvc_stmt *sql.Stmt
	psys_stmt *sql.Stmt
	pcst_stmt *sql.Stmt
	appLog    *log.Logger
	reqLog    *log.Logger
	errLog    *log.Logger
}

func (r *somaNodeReadHandler) run() {
	var err error

	for statement, prepStmt := range map[string]*sql.Stmt{
		stmt.ListNodes:       r.list_stmt,
		stmt.ShowNodes:       r.show_stmt,
		stmt.ShowConfigNodes: r.conf_stmt,
		stmt.SyncNodes:       r.sync_stmt,
		stmt.NodeOncProps:    r.ponc_stmt,
		stmt.NodeSvcProps:    r.psvc_stmt,
		stmt.NodeSysProps:    r.psys_stmt,
		stmt.NodeCstProps:    r.pcst_stmt,
	} {
		if prepStmt, err = r.conn.Prepare(statement); err != nil {
			r.errLog.Fatal(`node`, err, statement)
		}
		defer prepStmt.Close()
	}

runloop:
	for {
		select {
		case <-r.shutdown:
			break runloop
		case req := <-r.input:
			go func() {
				r.process(&req)
			}()
		}
	}
}

func (r *somaNodeReadHandler) process(q *somaNodeRequest) {
	var (
		nodeId, nodeName, nodeTeam, nodeServer, nodeState    string
		bucketId, repositoryId, instanceId, sourceInstanceId string
		view, oncallId, oncallName, serviceName, customId    string
		systemProp, value, customProp                        string
		nodeAsset                                            int
		nodeOnline, nodeDeleted                              bool
		rows                                                 *sql.Rows
		err                                                  error
		tx                                                   *sql.Tx
		checkConfigs                                         *[]proto.CheckConfig
	)
	result := somaResult{}

	switch q.action {
	case "list":
		r.reqLog.Printf("R: node/list")
		rows, err = r.list_stmt.Query()
		if result.SetRequestError(err) {
			goto dispatch
		}
		defer rows.Close()

		for rows.Next() {
			err = rows.Scan(&nodeId, &nodeName)
			result.Append(err, &somaNodeResult{
				Node: proto.Node{
					Id:   nodeId,
					Name: nodeName,
				},
			})
		}
	case `sync`:
		r.reqLog.Printf(`R: node/sync`)
		rows, err = r.sync_stmt.Query()
		if result.SetRequestError(err) {
			goto dispatch
		}
		defer rows.Close()

		for rows.Next() {
			err = rows.Scan(
				&nodeId,
				&nodeAsset,
				&nodeName,
				&nodeTeam,
				&nodeServer,
				&nodeOnline,
				&nodeDeleted,
			)
			result.Append(err, &somaNodeResult{
				Node: proto.Node{
					Id:        nodeId,
					AssetId:   uint64(nodeAsset),
					Name:      nodeName,
					TeamId:    nodeTeam,
					ServerId:  nodeServer,
					IsOnline:  nodeOnline,
					IsDeleted: nodeDeleted,
				},
			})
		}
	case "show":
		r.reqLog.Printf("R: node/show")
		err = r.show_stmt.QueryRow(q.Node.Id).Scan(
			&nodeId,
			&nodeAsset,
			&nodeName,
			&nodeTeam,
			&nodeServer,
			&nodeState,
			&nodeOnline,
			&nodeDeleted,
		)
		if err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			goto dispatch
		}
		node := proto.Node{
			Id:        nodeId,
			AssetId:   uint64(nodeAsset),
			Name:      nodeName,
			TeamId:    nodeTeam,
			ServerId:  nodeServer,
			State:     nodeState,
			IsOnline:  nodeOnline,
			IsDeleted: nodeDeleted,
		}

		// add configuration data
		err = r.conf_stmt.QueryRow(q.Node.Id).Scan(
			&nodeId,
			&nodeName,
			&bucketId,
			&repositoryId,
		)
		if err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			goto dispatch
		}
		node.Config = &proto.NodeConfig{
			RepositoryId: repositoryId,
			BucketId:     bucketId,
		}
		node.Properties = &[]proto.Property{}

		// oncall properties
		rows, err = r.ponc_stmt.Query(q.Node.Id)
		if result.SetRequestError(err) {
			goto dispatch
		}
		for rows.Next() {
			if err = rows.Scan(
				&instanceId,
				&sourceInstanceId,
				&view,
				&oncallId,
				&oncallName,
			); result.SetRequestError(err) {
				rows.Close()
				goto dispatch
			}
			*node.Properties = append(
				*node.Properties,
				proto.Property{
					Type:             `oncall`,
					RepositoryId:     repositoryId,
					BucketId:         bucketId,
					InstanceId:       instanceId,
					SourceInstanceId: sourceInstanceId,
					View:             view,
					Oncall: &proto.PropertyOncall{
						Id:   oncallId,
						Name: oncallName,
					},
				},
			)
		}
		if err = rows.Err(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}

		// service properties
		rows, err = r.psvc_stmt.Query(q.Node.Id)
		for rows.Next() {
			if err = rows.Scan(
				&instanceId,
				&sourceInstanceId,
				&view,
				&serviceName,
			); result.SetRequestError(err) {
				rows.Close()
				goto dispatch
			}
			*node.Properties = append(
				*node.Properties,
				proto.Property{
					Type:             `service`,
					RepositoryId:     repositoryId,
					BucketId:         bucketId,
					InstanceId:       instanceId,
					SourceInstanceId: sourceInstanceId,
					View:             view,
					Service: &proto.PropertyService{
						Name: serviceName,
					},
				},
			)
		}
		if err = rows.Err(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}

		// system properties
		rows, err = r.psys_stmt.Query(q.Node.Id)
		for rows.Next() {
			if err = rows.Scan(
				&instanceId,
				&sourceInstanceId,
				&view,
				&systemProp,
				&value,
			); result.SetRequestError(err) {
				rows.Close()
				goto dispatch
			}
			*node.Properties = append(
				*node.Properties,
				proto.Property{
					Type:             `system`,
					RepositoryId:     repositoryId,
					BucketId:         bucketId,
					InstanceId:       instanceId,
					SourceInstanceId: sourceInstanceId,
					View:             view,
					System: &proto.PropertySystem{
						Name:  systemProp,
						Value: value,
					},
				},
			)
		}
		if err = rows.Err(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}

		// custom properties
		rows, err = r.pcst_stmt.Query(q.Node.Id)
		for rows.Next() {
			if err = rows.Scan(
				&instanceId,
				&sourceInstanceId,
				&view,
				&customId,
				&value,
				&customProp,
			); result.SetRequestError(err) {
				rows.Close()
				goto dispatch
			}
			*node.Properties = append(
				*node.Properties,
				proto.Property{
					Type:             `custom`,
					RepositoryId:     repositoryId,
					BucketId:         bucketId,
					InstanceId:       instanceId,
					SourceInstanceId: sourceInstanceId,
					View:             view,
					Custom: &proto.PropertyCustom{
						Id:    customId,
						Name:  customProp,
						Value: value,
					},
				},
			)
		}
		if err = rows.Err(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}

		// add check configuration and instance information
		if tx, err = r.conn.Begin(); err != nil {
			result.SetRequestError(err)
			goto dispatch
		}
		checkConfigs, err = exportCheckConfigObjectTX(tx, q.Node.Id)
		if err != nil {
			tx.Rollback()
			result.SetRequestError(err)
			goto dispatch
		}
		if checkConfigs != nil && len(*checkConfigs) > 0 {
			node.Details = &proto.Details{
				CheckConfigs: checkConfigs,
			}
		}

		result.Append(err, &somaNodeResult{
			Node: node,
		})
	case "get_config":
		r.reqLog.Printf("R: node/get_config")
		err = r.conf_stmt.QueryRow(q.Node.Id).Scan(
			&nodeId,
			&nodeName,
			&bucketId,
			&repositoryId,
		)
		if err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			goto dispatch
		}

		result.Append(err, &somaNodeResult{
			Node: proto.Node{
				Id:   nodeId,
				Name: nodeName,
				Config: &proto.NodeConfig{
					RepositoryId: repositoryId,
					BucketId:     bucketId,
				},
			},
		})
	default:
		result.SetNotImplemented()
	}
dispatch:
	q.reply <- result
}

/* Write Access
 *
 */
type somaNodeWriteHandler struct {
	input    chan somaNodeRequest
	shutdown chan bool
	conn     *sql.DB
	add_stmt *sql.Stmt
	del_stmt *sql.Stmt
	prg_stmt *sql.Stmt
	upd_stmt *sql.Stmt
	appLog   *log.Logger
	reqLog   *log.Logger
	errLog   *log.Logger
}

func (w *somaNodeWriteHandler) run() {
	var err error

	for statement, prepStmt := range map[string]*sql.Stmt{
		stmt.NodeAdd:    w.add_stmt,
		stmt.NodeUpdate: w.upd_stmt,
		stmt.NodeDel:    w.del_stmt,
		stmt.NodePurge:  w.prg_stmt,
	} {
		if prepStmt, err = w.conn.Prepare(statement); err != nil {
			w.errLog.Fatal(`node`, err, statement)
		}
		defer prepStmt.Close()
	}

runloop:
	for {
		select {
		case <-w.shutdown:
			break runloop
		case req := <-w.input:
			go func() {
				w.process(&req)
			}()
		}
	}
}

func (w *somaNodeWriteHandler) process(q *somaNodeRequest) {
	var res sql.Result
	var err error
	result := somaResult{}

	switch q.action {
	case "add":
		w.reqLog.Printf("R: node/add for %s", q.Node.Name)
		id := uuid.NewV4()
		if q.Node.ServerId == "" {
			q.Node.ServerId = "00000000-0000-0000-0000-000000000000"
		}
		res, err = w.add_stmt.Exec(
			id.String(),
			q.Node.AssetId,
			q.Node.Name,
			q.Node.TeamId,
			q.Node.ServerId,
			q.Node.State,
			q.Node.IsOnline,
			false,
			q.user,
		)
		q.Node.Id = id.String()
	case `update`:
		w.reqLog.Printf("R: node/update for %s", q.Node.Id)
		res, err = w.upd_stmt.Exec(
			q.Node.AssetId,
			q.Node.Name,
			q.Node.TeamId,
			q.Node.ServerId,
			q.Node.IsOnline,
			q.Node.IsDeleted,
			q.Node.Id,
		)
		// TODO what has to be done for this undeployment?
	case "delete":
		w.reqLog.Printf("R: node/delete for %s", q.Node.Id)
		res, err = w.del_stmt.Exec(
			q.Node.Id,
		)
		// TODO trigger undeployment
	case "purge":
		w.reqLog.Printf("R: node/purge for %s", q.Node.Id)
		res, err = w.prg_stmt.Exec(
			q.Node.Id,
		)
	default:
		w.reqLog.Printf("R: unimplemented node/%s", q.action)
		result.SetNotImplemented()
		q.reply <- result
		return
	}
	if result.SetRequestError(err) {
		q.reply <- result
		return
	}

	rowCnt, _ := res.RowsAffected()
	switch {
	case rowCnt == 0:
		result.Append(errors.New("No rows affected"), &somaNodeResult{})
	case rowCnt > 1:
		result.Append(fmt.Errorf("Too many rows affected: %d", rowCnt),
			&somaNodeResult{})
	default:
		result.Append(nil, &somaNodeResult{
			Node: q.Node,
		})
	}
	q.reply <- result
}

/* Ops Access
 */
func (r *somaNodeReadHandler) shutdownNow() {
	r.shutdown <- true
}

func (w *somaNodeWriteHandler) shutdownNow() {
	w.shutdown <- true
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/rsa"
	"fmt"
	"net"
	"os"
	"path/filepath"

	"github.com/concourse/atc/atccmd"
	"github.com/concourse/flag"
	"github.com/concourse/tsa/tsacmd"
	flags "github.com/jessevdk/go-flags"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"

	"github.com/concourse/bin/bindata"
)

type WebCommand struct {
	*atccmd.ATCCommand

	*tsacmd.TSACommand `group:"TSA Configuration" namespace:"tsa"`
}

const cliArtifactsBindata = "cli-artifacts"

func (WebCommand) lessenRequirements(command *flags.Command) {
	// defaults to address from external URL
	command.FindOptionByLongName("tsa-peer-ip").Required = false

	// defaults to atc external URL
	command.FindOptionByLongName("tsa-atc-url").Required = false

	// defaults to atc session signing key
	command.FindOptionByLongName("tsa-session-signing-key").Required = false
}

func (cmd *WebCommand) Execute(args []string) error {
	runner, err := cmd.Runner(args)
	if err != nil {
		return err
	}

	return <-ifrit.Invoke(sigmon.New(runner)).Wait()
}

func (cmd *WebCommand) Runner(args []string) (ifrit.Runner, error) {
	err := bindata.RestoreAssets(os.TempDir(), cliArtifactsBindata)
	if err == nil {
		cmd.ATCCommand.CLIArtifactsDir = flag.Dir(filepath.Join(os.TempDir(), cliArtifactsBindata))
	}

	cmd.populateTSAFlagsFromATCFlags()

	atcRunner, err := cmd.ATCCommand.Runner(args)
	if err != nil {
		return nil, err
	}

	tsaRunner, err := cmd.TSACommand.Runner(args)
	if err != nil {
		return nil, err
	}

	return grouper.NewParallel(os.Interrupt, grouper.Members{
		{Name: "atc", Runner: atcRunner},
		{Name: "tsa", Runner: tsaRunner},
	}), nil
}

func (cmd *WebCommand) populateTSAFlagsFromATCFlags() error {
	cmd.TSACommand.SessionSigningKey = cmd.ATCCommand.Auth.AuthFlags.SigningKey

	if cmd.ATCCommand.Auth.AuthFlags.SigningKey.PrivateKey == nil &&
		cmd.TSACommand.SessionSigningKey.PrivateKey == nil {
		signingKey, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return fmt.Errorf("failed to generate session signing key: %s", err)
		}

		cmd.ATCCommand.Auth.AuthFlags.SigningKey = &flag.PrivateKey{PrivateKey: signingKey}
		cmd.TSACommand.SessionSigningKey = &flag.PrivateKey{PrivateKey: signingKey}
	}

	if len(cmd.TSACommand.ATCURLs) == 0 {
		cmd.TSACommand.ATCURLs = []flag.URL{cmd.ATCCommand.PeerURL}
	}

	host, _, err := net.SplitHostPort(cmd.ATCCommand.PeerURL.URL.Host)
	if err != nil {
		return err
	}

	cmd.TSACommand.PeerIP = host

	cmd.TSACommand.Metrics.YellerAPIKey = cmd.ATCCommand.Metrics.YellerAPIKey
	cmd.TSACommand.Metrics.YellerEnvironment = cmd.ATCCommand.Metrics.YellerEnvironment

	return nil
}
<commit_msg>The web command no longer brings up the TSA when running migration commands<commit_after>package main

import (
	"crypto/rand"
	"crypto/rsa"
	"fmt"
	"net"
	"os"
	"path/filepath"

	"github.com/concourse/atc/atccmd"
	"github.com/concourse/flag"
	"github.com/concourse/tsa/tsacmd"
	flags "github.com/jessevdk/go-flags"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/sigmon"

	"github.com/concourse/bin/bindata"
)

type WebCommand struct {
	*atccmd.ATCCommand

	*tsacmd.TSACommand `group:"TSA Configuration" namespace:"tsa"`
}

const cliArtifactsBindata = "cli-artifacts"

func (WebCommand) lessenRequirements(command *flags.Command) {
	// defaults to address from external URL
	command.FindOptionByLongName("tsa-peer-ip").Required = false

	// defaults to atc external URL
	command.FindOptionByLongName("tsa-atc-url").Required = false

	// defaults to atc session signing key
	command.FindOptionByLongName("tsa-session-signing-key").Required = false
}

func (cmd *WebCommand) Execute(args []string) error {
	runner, err := cmd.Runner(args)
	if err != nil {
		return err
	}

	return <-ifrit.Invoke(sigmon.New(runner)).Wait()
}

func (cmd *WebCommand) Runner(args []string) (ifrit.Runner, error) {
	err := bindata.RestoreAssets(os.TempDir(), cliArtifactsBindata)
	if err == nil {
		cmd.ATCCommand.CLIArtifactsDir = flag.Dir(filepath.Join(os.TempDir(), cliArtifactsBindata))
	}

	cmd.populateTSAFlagsFromATCFlags()

	atcRunner, shouldSkipTSA, err := cmd.ATCCommand.Runner(args)
	if err != nil {
		return nil, err
	}

	if shouldSkipTSA {
		return atcRunner, nil
	}

	tsaRunner, err := cmd.TSACommand.Runner(args)
	if err != nil {
		return nil, err
	}

	return grouper.NewParallel(os.Interrupt, grouper.Members{
		{Name: "atc", Runner: atcRunner},
		{Name: "tsa", Runner: tsaRunner},
	}), nil
}

func (cmd *WebCommand) populateTSAFlagsFromATCFlags() error {
	cmd.TSACommand.SessionSigningKey = cmd.ATCCommand.Auth.AuthFlags.SigningKey

	if cmd.ATCCommand.Auth.AuthFlags.SigningKey.PrivateKey == nil &&
		cmd.TSACommand.SessionSigningKey.PrivateKey == nil {
		signingKey, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return fmt.Errorf("failed to generate session signing key: %s", err)
		}

		cmd.ATCCommand.Auth.AuthFlags.SigningKey = &flag.PrivateKey{PrivateKey: signingKey}
		cmd.TSACommand.SessionSigningKey = &flag.PrivateKey{PrivateKey: signingKey}
	}

	if len(cmd.TSACommand.ATCURLs) == 0 {
		cmd.TSACommand.ATCURLs = []flag.URL{cmd.ATCCommand.PeerURL}
	}

	host, _, err := net.SplitHostPort(cmd.ATCCommand.PeerURL.URL.Host)
	if err != nil {
		return err
	}

	cmd.TSACommand.PeerIP = host

	cmd.TSACommand.Metrics.YellerAPIKey = cmd.ATCCommand.Metrics.YellerAPIKey
	cmd.TSACommand.Metrics.YellerEnvironment = cmd.ATCCommand.Metrics.YellerEnvironment

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/gempir/gempbot/pkg/config"
	"github.com/gempir/gempbot/pkg/eventsub"
	"github.com/gempir/gempbot/pkg/helix"
	"github.com/gempir/gempbot/pkg/log"
	"github.com/gempir/gempbot/pkg/store"
	nickHelix "github.com/nicklaw5/helix/v2"
)

var (
	cfg         *config.Config
	db          *store.Database
	helixClient *helix.Client
)

func main() {
	cfg = config.FromEnv()
	db = store.NewDatabase(cfg)
	helixClient = helix.NewClient(cfg, db)
	subscriptionManager := eventsub.NewSubscriptionManager(cfg, db, helixClient)

	for _, sub := range db.GetAllSubscriptions() {
		if sub.Type == nickHelix.EventSubTypeChannelPointsCustomRewardAdd {
			log.Infof("sub: %s foreignId: '%v'", sub.SubscriptionID, sub.ForeignID)

			if sub.ForeignID == "" {
				log.Warnf("Found subscription reward without foreignID channel: %s id: %s", sub.TargetTwitchID, sub.SubscriptionID)
				err := subscriptionManager.RemoveSubscription(sub.SubscriptionID)
				if err != nil {
					log.Error(err)
				}
			}
		}
	}
}
<commit_msg>more logging<commit_after>package main

import (
	"github.com/gempir/gempbot/pkg/config"
	"github.com/gempir/gempbot/pkg/eventsub"
	"github.com/gempir/gempbot/pkg/helix"
	"github.com/gempir/gempbot/pkg/log"
	"github.com/gempir/gempbot/pkg/store"
	nickHelix "github.com/nicklaw5/helix/v2"
)

var (
	cfg         *config.Config
	db          *store.Database
	helixClient *helix.Client
)

func main() {
	cfg = config.FromEnv()
	db = store.NewDatabase(cfg)
	helixClient = helix.NewClient(cfg, db)
	subscriptionManager := eventsub.NewSubscriptionManager(cfg, db, helixClient)

	for _, sub := range db.GetAllSubscriptions() {
		log.Infof("type: '%v'", sub.Type)
		if sub.Type == nickHelix.EventSubTypeChannelPointsCustomRewardAdd {
			log.Infof("sub: %s foreignId: '%v'", sub.SubscriptionID, sub.ForeignID)

			if sub.ForeignID == "" {
				log.Warnf("Found subscription reward without foreignID channel: %s id: %s", sub.TargetTwitchID, sub.SubscriptionID)
				err := subscriptionManager.RemoveSubscription(sub.SubscriptionID)
				if err != nil {
					log.Error(err)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/cerana/cerana/acomm"
)

type statsPusher struct {
	config   *config
	tracker  *acomm.Tracker
	stopChan chan struct{}
	wg       sync.WaitGroup
}

func newStatsPusher(c *config) (*statsPusher, error) {
	tracker, err := acomm.NewTracker("", nil, nil, c.requestTimeout())
	if err != nil {
		return nil, err
	}

	return &statsPusher{
		config:   c,
		tracker:  tracker,
		stopChan: make(chan struct{}),
	}, nil
}

func (s *statsPusher) run() error {
	if err := s.tracker.Start(); err != nil {
		return err
	}
	s.startHeartbeat("node", s.nodeHeartbeat, s.config.nodeTTL())
	s.startHeartbeat("dataset", s.datasetHeartbeats, s.config.datasetTTL())
	s.startHeartbeat("bundle", s.bundleHeartbeats, s.config.bundleTTL())
	return nil
}

func (s *statsPusher) startHeartbeat(name string, fn func() error, desiredInterval time.Duration) {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		lastStart := time.Time{}
		for {
			// Try to start fn at the desired time interval, accounting for
			// time time it takes for fn to complete while also preventing fn
			// from running more than once at a time.
			interval := desiredInterval
			since := time.Since(lastStart)
			if since < interval {
				interval = since
			}
			timeChan := time.After(interval)

			select {
			case <-s.stopChan:
				break
			case lastStart = <-timeChan:
				if err := fn(); err != nil {
					logrus.WithFields(logrus.Fields{
						"name":  name,
						"error": err,
					}).Error("error while running heartbeat")
				}
			}
		}
	}()
}

func (s *statsPusher) stop() {
	close(s.stopChan)
	s.wg.Wait()

	s.tracker.Stop()
}

func (s *statsPusher) stopOnSignal(signals ...os.Signal) {
	if len(signals) == 0 {
		signals = []os.Signal{os.Interrupt, os.Kill, syscall.SIGTERM}
	}

	sigChan := make(chan os.Signal)
	signal.Notify(sigChan, signals...)
	sig := <-sigChan
	logrus.WithFields(logrus.Fields{
		"signal": sig,
	}).Info("signal received, stopping")

	s.stop()
}
<commit_msg>Comment typo and simplification in statspusher<commit_after>package main

import (
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/cerana/cerana/acomm"
)

type statsPusher struct {
	config   *config
	tracker  *acomm.Tracker
	stopChan chan struct{}
	wg       sync.WaitGroup
}

func newStatsPusher(c *config) (*statsPusher, error) {
	tracker, err := acomm.NewTracker("", nil, nil, c.requestTimeout())
	if err != nil {
		return nil, err
	}

	return &statsPusher{
		config:   c,
		tracker:  tracker,
		stopChan: make(chan struct{}),
	}, nil
}

func (s *statsPusher) run() error {
	if err := s.tracker.Start(); err != nil {
		return err
	}
	s.startHeartbeat("node", s.nodeHeartbeat, s.config.nodeTTL())
	s.startHeartbeat("dataset", s.datasetHeartbeats, s.config.datasetTTL())
	s.startHeartbeat("bundle", s.bundleHeartbeats, s.config.bundleTTL())
	return nil
}

func (s *statsPusher) startHeartbeat(name string, fn func() error, desiredInterval time.Duration) {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		lastStart := time.Time{}
		for {
			// Try to start fn at the desired time interval, accounting for
			// the time it takes for fn to complete while also preventing fn
			// from running more than once at a time.
			interval := desiredInterval
			since := time.Since(lastStart)
			if since < interval {
				interval = since
			}

			select {
			case <-s.stopChan:
				break
			case lastStart = <-time.After(interval):
				if err := fn(); err != nil {
					logrus.WithFields(logrus.Fields{
						"name":  name,
						"error": err,
					}).Error("error while running heartbeat")
				}
			}
		}
	}()
}

func (s *statsPusher) stop() {
	close(s.stopChan)
	s.wg.Wait()

	s.tracker.Stop()
}

func (s *statsPusher) stopOnSignal(signals ...os.Signal) {
	if len(signals) == 0 {
		signals = []os.Signal{os.Interrupt, os.Kill, syscall.SIGTERM}
	}

	sigChan := make(chan os.Signal)
	signal.Notify(sigChan, signals...)
	sig := <-sigChan
	logrus.WithFields(logrus.Fields{
		"signal": sig,
	}).Info("signal received, stopping")

	s.stop()
}
<|endoftext|>


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

// kexec executes a new kernel over the running kernel (u-root).
//
// Synopsis:
//     kexec [--initrd=FILE] [--command-line=STRING] [-l] [-e] [KERNELIMAGE]
//
// Description:
//		 Loads a kernel for later execution.
//
// Options:
//      --append string        Append to the kernel command line
//  -c, --cmdline string       Append to the kernel command line
//  -d, --debug                Print debug info (default true)
//  -e, --exec                 Execute a currently loaded kernel
//  -x, --extra string         Add a cpio containing extra files
//      --initramfs string     Use file as the kernel's initial ramdisk
//  -i, --initrd string        Use file as the kernel's initial ramdisk
//  -l, --load                 Load the new kernel into the current kernel
//  -L, --loadsyscall          Use the kexec load syscall (not file_load) (default true)
//      --module stringArray   Load multiboot module with command line args (e.g --module="mod arg1")
//  -p, --purgatory string     pick a purgatory, use '-p xyz' to get a list (default "default")
//      --reuse-cmdline        Use the kernel command line from running system

package main

import (
	"bytes"
	"io"
	"log"
	"os"
	"strings"

	flag "github.com/spf13/pflag"

	"github.com/u-root/u-root/pkg/boot"
	"github.com/u-root/u-root/pkg/boot/kexec"
	"github.com/u-root/u-root/pkg/boot/multiboot"
	"github.com/u-root/u-root/pkg/cmdline"
	"github.com/u-root/u-root/pkg/cpio"
	"github.com/u-root/u-root/pkg/uio"
)

type options struct {
	loadSyscall  bool
	cmdline      string
	reuseCmdline bool
	initramfs    string
	load         bool
	exec         bool
	debug        bool
	extra        string
	purgatory    string
	modules      []string
}

func registerFlags() *options {
	o := &options{}
	flag.BoolVarP(&o.loadSyscall, "loadsyscall", "L", false, "Use the kexec load syscall (not file_load)")
	flag.StringVarP(&o.cmdline, "cmdline", "c", "", "Append to the kernel command line")
	flag.StringVar(&o.cmdline, "append", "", "Append to the kernel command line")
	flag.StringVarP(&o.extra, "extra", "x", "", "Add a cpio containing extra files")
	flag.BoolVar(&o.reuseCmdline, "reuse-cmdline", false, "Use the kernel command line from running system")
	flag.StringVarP(&o.initramfs, "initrd", "i", "", "Use file as the kernel's initial ramdisk")
	flag.StringVar(&o.initramfs, "initramfs", "", "Use file as the kernel's initial ramdisk")
	flag.BoolVarP(&o.load, "load", "l", false, "Load the new kernel into the current kernel")
	flag.BoolVarP(&o.exec, "exec", "e", false, "Execute a currently loaded kernel")
	flag.BoolVarP(&o.debug, "debug", "d", false, "Print debug info")
	flag.StringArrayVar(&o.modules, "module", nil, `Load multiboot module with command line args (e.g --module="mod arg1")`)

	// This is broken out as it is almost never to be used. But it is valueable, nonetheless.
	flag.StringVarP(&o.purgatory, "purgatory", "p", "default", "pick a purgatory, use '-p xyz' to get a list")
	return o
}

func main() {
	opts := registerFlags()
	flag.Parse()

	if opts.debug {
		kexec.Debug = log.Printf
	}

	if (!opts.exec && flag.NArg() == 0) || flag.NArg() > 1 {
		flag.PrintDefaults()
		log.Fatalf("usage: kexec [flags] kernelname OR kexec -e")
	}

	if opts.cmdline != "" && opts.reuseCmdline {
		flag.PrintDefaults()
		log.Fatalf("--reuse-cmdline and other command line options are mutually exclusive")
	}

	if !opts.load && !opts.exec {
		opts.load = true
		opts.exec = true
	}

	newCmdline := opts.cmdline
	if opts.reuseCmdline {
		procCmdLine := cmdline.NewCmdLine()
		if procCmdLine.Err != nil {
			log.Fatal("Couldn't read /proc/cmdline")
		} else {
			newCmdline = procCmdLine.Raw
		}
	}

	if err := kexec.SelectPurgator(opts.purgatory); err != nil {
		log.Fatal(err)
	}
	if opts.load {
		kernelpath := flag.Arg(0)
		kernel, err := os.Open(kernelpath)
		if err != nil {
			log.Fatal(err)
		}
		defer kernel.Close()
		var image boot.OSImage
		if err := multiboot.Probe(kernel); err == nil {
			image = &boot.MultibootImage{
				Modules: multiboot.LazyOpenModules(opts.modules),
				Kernel:  kernel,
				Cmdline: newCmdline,
			}
		} else {
			var files []io.ReaderAt
			if len(opts.extra) > 0 {
				b := &bytes.Buffer{}
				archiver, err := cpio.Format("newc")
				if err != nil {
					log.Fatal(err)
				}
				w := archiver.Writer(b)
				cr := cpio.NewRecorder()
				// to deconflict names, we may want to prepend the names with
				// kexec_extra/ or something.
				for _, n := range strings.Fields(opts.extra) {
					rec, err := cr.GetRecord(n)
					if err != nil {
						log.Fatalf("Getting record of %q failed: %v", n, err)
					}
					if err := w.WriteRecord(rec); err != nil {
						log.Fatalf("Writing record %q failed: %v", n, err)
					}
				}
				if err := cpio.WriteTrailer(w); err != nil {
					log.Fatalf("Error writing trailer record: %v", err)
				}
				files = append(files, bytes.NewReader(b.Bytes()))
			}
			if opts.initramfs != "" {
				for _, n := range strings.Fields(opts.initramfs) {
					files = append(files, uio.NewLazyFile(n))
				}
			}
			var i io.ReaderAt
			if len(files) > 0 {
				i = boot.CatInitrds(files...)
			}

			image = &boot.LinuxImage{
				Kernel:      uio.NewLazyFile(kernelpath),
				Initrd:      i,
				Cmdline:     newCmdline,
				LoadSyscall: opts.loadSyscall,
			}
		}
		if err := image.Load(opts.debug); err != nil {
			log.Fatal(err)
		}
	}

	if opts.exec {
		if err := kexec.Reboot(); err != nil {
			log.Fatalf("%v", err)
		}
	}
}
<commit_msg>cmds/core/kexec: re-add the package comment<commit_after>// Copyright 2015-2018 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// kexec executes a new kernel over the running kernel (u-root).
//
// Synopsis:
//     kexec [--initrd=FILE] [--command-line=STRING] [-l] [-e] [KERNELIMAGE]
//
// Description:
//		 Loads a kernel for later execution.
//
// Options:
//      --append string        Append to the kernel command line
//  -c, --cmdline string       Append to the kernel command line
//  -d, --debug                Print debug info (default true)
//  -e, --exec                 Execute a currently loaded kernel
//  -x, --extra string         Add a cpio containing extra files
//      --initramfs string     Use file as the kernel's initial ramdisk
//  -i, --initrd string        Use file as the kernel's initial ramdisk
//  -l, --load                 Load the new kernel into the current kernel
//  -L, --loadsyscall          Use the kexec load syscall (not file_load) (default true)
//      --module stringArray   Load multiboot module with command line args (e.g --module="mod arg1")
//  -p, --purgatory string     pick a purgatory, use '-p xyz' to get a list (default "default")
//      --reuse-cmdline        Use the kernel command line from running system
package main

import (
	"bytes"
	"io"
	"log"
	"os"
	"strings"

	flag "github.com/spf13/pflag"

	"github.com/u-root/u-root/pkg/boot"
	"github.com/u-root/u-root/pkg/boot/kexec"
	"github.com/u-root/u-root/pkg/boot/multiboot"
	"github.com/u-root/u-root/pkg/cmdline"
	"github.com/u-root/u-root/pkg/cpio"
	"github.com/u-root/u-root/pkg/uio"
)

type options struct {
	loadSyscall  bool
	cmdline      string
	reuseCmdline bool
	initramfs    string
	load         bool
	exec         bool
	debug        bool
	extra        string
	purgatory    string
	modules      []string
}

func registerFlags() *options {
	o := &options{}
	flag.BoolVarP(&o.loadSyscall, "loadsyscall", "L", false, "Use the kexec load syscall (not file_load)")
	flag.StringVarP(&o.cmdline, "cmdline", "c", "", "Append to the kernel command line")
	flag.StringVar(&o.cmdline, "append", "", "Append to the kernel command line")
	flag.StringVarP(&o.extra, "extra", "x", "", "Add a cpio containing extra files")
	flag.BoolVar(&o.reuseCmdline, "reuse-cmdline", false, "Use the kernel command line from running system")
	flag.StringVarP(&o.initramfs, "initrd", "i", "", "Use file as the kernel's initial ramdisk")
	flag.StringVar(&o.initramfs, "initramfs", "", "Use file as the kernel's initial ramdisk")
	flag.BoolVarP(&o.load, "load", "l", false, "Load the new kernel into the current kernel")
	flag.BoolVarP(&o.exec, "exec", "e", false, "Execute a currently loaded kernel")
	flag.BoolVarP(&o.debug, "debug", "d", false, "Print debug info")
	flag.StringArrayVar(&o.modules, "module", nil, `Load multiboot module with command line args (e.g --module="mod arg1")`)

	// This is broken out as it is almost never to be used. But it is valueable, nonetheless.
	flag.StringVarP(&o.purgatory, "purgatory", "p", "default", "pick a purgatory, use '-p xyz' to get a list")
	return o
}

func main() {
	opts := registerFlags()
	flag.Parse()

	if opts.debug {
		kexec.Debug = log.Printf
	}

	if (!opts.exec && flag.NArg() == 0) || flag.NArg() > 1 {
		flag.PrintDefaults()
		log.Fatalf("usage: kexec [flags] kernelname OR kexec -e")
	}

	if opts.cmdline != "" && opts.reuseCmdline {
		flag.PrintDefaults()
		log.Fatalf("--reuse-cmdline and other command line options are mutually exclusive")
	}

	if !opts.load && !opts.exec {
		opts.load = true
		opts.exec = true
	}

	newCmdline := opts.cmdline
	if opts.reuseCmdline {
		procCmdLine := cmdline.NewCmdLine()
		if procCmdLine.Err != nil {
			log.Fatal("Couldn't read /proc/cmdline")
		} else {
			newCmdline = procCmdLine.Raw
		}
	}

	if err := kexec.SelectPurgator(opts.purgatory); err != nil {
		log.Fatal(err)
	}
	if opts.load {
		kernelpath := flag.Arg(0)
		kernel, err := os.Open(kernelpath)
		if err != nil {
			log.Fatal(err)
		}
		defer kernel.Close()
		var image boot.OSImage
		if err := multiboot.Probe(kernel); err == nil {
			image = &boot.MultibootImage{
				Modules: multiboot.LazyOpenModules(opts.modules),
				Kernel:  kernel,
				Cmdline: newCmdline,
			}
		} else {
			var files []io.ReaderAt
			if len(opts.extra) > 0 {
				b := &bytes.Buffer{}
				archiver, err := cpio.Format("newc")
				if err != nil {
					log.Fatal(err)
				}
				w := archiver.Writer(b)
				cr := cpio.NewRecorder()
				// to deconflict names, we may want to prepend the names with
				// kexec_extra/ or something.
				for _, n := range strings.Fields(opts.extra) {
					rec, err := cr.GetRecord(n)
					if err != nil {
						log.Fatalf("Getting record of %q failed: %v", n, err)
					}
					if err := w.WriteRecord(rec); err != nil {
						log.Fatalf("Writing record %q failed: %v", n, err)
					}
				}
				if err := cpio.WriteTrailer(w); err != nil {
					log.Fatalf("Error writing trailer record: %v", err)
				}
				files = append(files, bytes.NewReader(b.Bytes()))
			}
			if opts.initramfs != "" {
				for _, n := range strings.Fields(opts.initramfs) {
					files = append(files, uio.NewLazyFile(n))
				}
			}
			var i io.ReaderAt
			if len(files) > 0 {
				i = boot.CatInitrds(files...)
			}

			image = &boot.LinuxImage{
				Kernel:      uio.NewLazyFile(kernelpath),
				Initrd:      i,
				Cmdline:     newCmdline,
				LoadSyscall: opts.loadSyscall,
			}
		}
		if err := image.Load(opts.debug); err != nil {
			log.Fatal(err)
		}
	}

	if opts.exec {
		if err := kexec.Reboot(); err != nil {
			log.Fatalf("%v", err)
		}
	}
}
<|endoftext|>