<commit_before>package main

import (
	"exercism"
	"fmt"
	"github.com/codegangsta/cli"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	app := cli.NewApp()
	app.Name = "exercism"
	app.Usage = "A command line tool to interact with http://exercism.io"
	app.Version = exercism.VERSION
	app.Commands = []cli.Command{
		{
			Name:      "demo",
			ShortName: "d",
			Usage:     "Fetch first assignment for each language from exercism.io",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					demoDir, err2 := exercism.DemoDirectory()
					if err2 != nil {
						err = err2
						fmt.Println(err)
						return
					}
					config = exercism.Config{
						Hostname: "http://exercism.io",
						ApiKey: "",
						ExercismDirectory: demoDir,
					}
				}
				assignments, err := exercism.FetchAssignments(config.Hostname,
					exercism.FetchEndpoints["demo"], config.ApiKey)
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := exercism.SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "fetch",
			ShortName: "f",
			Usage:     "Fetch current assignment from exercism.io",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				assignments, err := exercism.FetchAssignments(config.Hostname,
					exercism.FetchEndpoints["current"], config.ApiKey)
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := exercism.SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "login",
			ShortName: "l",
			Usage:     "Save exercism.io api credentials",
			Action: func(c *cli.Context) {
				exercism.ConfigToFile(exercism.HomeDir(), askForConfigInfo())
			},
		},
		{
			Name:      "logout",
			ShortName: "o",
			Usage:     "Clear exercism.io api credentials",
			Action: func(c *cli.Context) {
				exercism.Logout(exercism.HomeDir())
			},
		},
		{
			Name:      "peek",
			ShortName: "p",
			Usage:     "Fetch upcoming assignment from exercism.io",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				assignments, err := exercism.FetchAssignments(config.Hostname,
					exercism.FetchEndpoints["next"], config.ApiKey)
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := exercism.SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "submit",
			ShortName: "s",
			Usage:     "Submit code to exercism.io on your current assignment",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				if len(c.Args()) == 0 {
					fmt.Println("Please enter a file name")
					return
				}

				filename := c.Args()[0]

				// Make filename relative to config.ExercismDirectory.
				absPath, err := absolutePath(filename)
				if err != nil {
					fmt.Printf("Couldn't find %v: %v\n", filename, err)
					return
				}
				exDir := config.ExercismDirectory + string(filepath.Separator)
				if !strings.HasPrefix(absPath, exDir) {
					fmt.Printf("%v is not under your exercism project path (%v)\n", absPath, exDir)
					return
				}
				filename = absPath[len(exDir):]

				if exercism.IsTest(filename) {
					fmt.Println("It looks like this is a test, please enter an example file name.")
					return
				}

				code, err := ioutil.ReadFile(absPath)
				if err != nil {
					fmt.Printf("Error reading %v: %v\n", absPath, err)
					return
				}

				response, err := exercism.SubmitAssignment(config.Hostname, config.ApiKey, filename, code)
				if err != nil {
					fmt.Printf("There was an issue with your submission: %v\n", err)
					return
				}

				fmt.Printf("For feedback on your submission visit %s%s.\n",
					config.Hostname, response.SubmissionPath)

			},
		},
		{
			Name:      "whoami",
			ShortName: "w",
			Usage:     "Get the github username that you are logged in as",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				fmt.Println(config.GithubUsername)
			},
		},
	}
	app.Run(os.Args)
}

func askForConfigInfo() (c exercism.Config) {
	var un, key, dir string

	currentDir, err := os.Getwd()
	if err != nil {
		panic(err)
	}

	fmt.Print("Your GitHub username: ")
	_, err = fmt.Scanln(&un)
	if err != nil {
		panic(err)
	}

	fmt.Print("Your exercism.io API key: ")
	_, err = fmt.Scanln(&key)
	if err != nil {
		panic(err)
	}

	fmt.Println("What is your exercism exercises project path?")
	fmt.Printf("Press Enter to select the default (%s):\n", currentDir)
	fmt.Print("> ")
	_, err = fmt.Scanln(&dir)
	if err != nil && err.Error() != "unexpected newline" {
		panic(err)
	}
	dir, err = absolutePath(dir)
	if err != nil {
		panic(err)
	}

	if dir == "" {
		dir = currentDir
	}

	return exercism.Config{un, key, exercism.ReplaceTilde(dir), "http://exercism.io"}
}

func absolutePath(path string) (string, error) {
	path, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	return filepath.EvalSymlinks(path)
}
<commit_msg>Refactor error handling in demo command<commit_after>package main

import (
	"exercism"
	"fmt"
	"github.com/codegangsta/cli"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	app := cli.NewApp()
	app.Name = "exercism"
	app.Usage = "A command line tool to interact with http://exercism.io"
	app.Version = exercism.VERSION
	app.Commands = []cli.Command{
		{
			Name:      "demo",
			ShortName: "d",
			Usage:     "Fetch first assignment for each language from exercism.io",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					demoDir, err := exercism.DemoDirectory()
					if err != nil {
						fmt.Println(err)
						return
					}
					config = exercism.Config{
						Hostname: "http://exercism.io",
						ApiKey: "",
						ExercismDirectory: demoDir,
					}
				}
				assignments, err := exercism.FetchAssignments(config.Hostname,
					exercism.FetchEndpoints["demo"], config.ApiKey)
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := exercism.SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "fetch",
			ShortName: "f",
			Usage:     "Fetch current assignment from exercism.io",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				assignments, err := exercism.FetchAssignments(config.Hostname,
					exercism.FetchEndpoints["current"], config.ApiKey)
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := exercism.SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "login",
			ShortName: "l",
			Usage:     "Save exercism.io api credentials",
			Action: func(c *cli.Context) {
				exercism.ConfigToFile(exercism.HomeDir(), askForConfigInfo())
			},
		},
		{
			Name:      "logout",
			ShortName: "o",
			Usage:     "Clear exercism.io api credentials",
			Action: func(c *cli.Context) {
				exercism.Logout(exercism.HomeDir())
			},
		},
		{
			Name:      "peek",
			ShortName: "p",
			Usage:     "Fetch upcoming assignment from exercism.io",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				assignments, err := exercism.FetchAssignments(config.Hostname,
					exercism.FetchEndpoints["next"], config.ApiKey)
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := exercism.SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "submit",
			ShortName: "s",
			Usage:     "Submit code to exercism.io on your current assignment",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				if len(c.Args()) == 0 {
					fmt.Println("Please enter a file name")
					return
				}

				filename := c.Args()[0]

				// Make filename relative to config.ExercismDirectory.
				absPath, err := absolutePath(filename)
				if err != nil {
					fmt.Printf("Couldn't find %v: %v\n", filename, err)
					return
				}
				exDir := config.ExercismDirectory + string(filepath.Separator)
				if !strings.HasPrefix(absPath, exDir) {
					fmt.Printf("%v is not under your exercism project path (%v)\n", absPath, exDir)
					return
				}
				filename = absPath[len(exDir):]

				if exercism.IsTest(filename) {
					fmt.Println("It looks like this is a test, please enter an example file name.")
					return
				}

				code, err := ioutil.ReadFile(absPath)
				if err != nil {
					fmt.Printf("Error reading %v: %v\n", absPath, err)
					return
				}

				response, err := exercism.SubmitAssignment(config.Hostname, config.ApiKey, filename, code)
				if err != nil {
					fmt.Printf("There was an issue with your submission: %v\n", err)
					return
				}

				fmt.Printf("For feedback on your submission visit %s%s.\n",
					config.Hostname, response.SubmissionPath)

			},
		},
		{
			Name:      "whoami",
			ShortName: "w",
			Usage:     "Get the github username that you are logged in as",
			Action: func(c *cli.Context) {
				config, err := exercism.ConfigFromFile(exercism.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				fmt.Println(config.GithubUsername)
			},
		},
	}
	app.Run(os.Args)
}

func askForConfigInfo() (c exercism.Config) {
	var un, key, dir string

	currentDir, err := os.Getwd()
	if err != nil {
		panic(err)
	}

	fmt.Print("Your GitHub username: ")
	_, err = fmt.Scanln(&un)
	if err != nil {
		panic(err)
	}

	fmt.Print("Your exercism.io API key: ")
	_, err = fmt.Scanln(&key)
	if err != nil {
		panic(err)
	}

	fmt.Println("What is your exercism exercises project path?")
	fmt.Printf("Press Enter to select the default (%s):\n", currentDir)
	fmt.Print("> ")
	_, err = fmt.Scanln(&dir)
	if err != nil && err.Error() != "unexpected newline" {
		panic(err)
	}
	dir, err = absolutePath(dir)
	if err != nil {
		panic(err)
	}

	if dir == "" {
		dir = currentDir
	}

	return exercism.Config{un, key, exercism.ReplaceTilde(dir), "http://exercism.io"}
}

func absolutePath(path string) (string, error) {
	path, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	return filepath.EvalSymlinks(path)
}
<|endoftext|>


<commit_before>// Copyright 2016 FullStory, Inc.
//
// 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 smservice

import (
	"fmt"
	"runtime"
	"sort"
	"strings"
	"time"

	"github.com/fullstorydev/gosolr/smutil"
	"github.com/fullstorydev/gosolr/solrcheck"
	"github.com/fullstorydev/gosolr/solrman/smmodel"
	"github.com/fullstorydev/gosolr/solrman/solrmanapi"
	"github.com/samuel/go-zookeeper/zk"
)

const (
	movesPerCycle                 = 10               // How many shards to move at a time
	iterationSleep                = 1 * time.Minute  // How long to sleep each attempt, no matter what
	quiescenceSleep               = 10 * time.Minute // How long to sleep when stability is reached
	splitsPerMachine              = 4                // How many shard splitting jobs to schedule on 1 physical machine at a time
	splitShardsWithDocCount       = 4000000          // Split shards with doc count > this
	allSplitsDocCountTrigger      = 4004000          // But don't do any splits until at least one shard > this
	allowedMinToMaxShardSizeRatio = 0.2              // Ratio of smallest shard to biggest shard < this then warn about imbalance
	maxShardsPerMachine           = 16               // Maximum number of shards per machine in the cluster
)

// Runs the main solr management loop, never returns.
func (s *SolrManService) RunSolrMan() {
	s.setStatusOp("solrman is starting up")
	clusterStateGolden := true // assume true to start

	first := true
	for {
		if !first {
			time.Sleep(iterationSleep)
		}
		first = false

		if s.ZooClient.State() != zk.StateHasSession {
			s.setStatusOp("not connected to zk")
			s.Logger.Warningf("not connected to zk")
			continue
		}

		if s.hasInProgressOps() {
			s.clearStatusOp()
			s.Logger.Debugf("in progress ops")
			continue
		}

		if s.Storage.IsDisabled() {
			s.setStatusOp("solrman is disabled")
			s.Logger.Infof("solrman is disabled")
			continue
		}

		evacuatingNodes, err := s.Storage.GetEvacuateNodeList()
		if err != nil {
			s.setStatusOp("failed to determine hosts to evacuate")
			s.Logger.Errorf("failed to determine hosts to evacuate: %s", err)
			continue
		}

		clusterState, err := s.SolrMonitor.GetCurrentState()
		if err != nil {
			s.setStatusOp("failed to retrieve cluster state")
			s.Logger.Errorf("failed to retrieve cluster state: %s", err)
			continue
		}
		liveNodes, err := s.SolrMonitor.GetLiveNodes()
		if err != nil {
			s.setStatusOp("failed to retrieve live nodes")
			s.Logger.Errorf("failed to retrieve live nodes: %s", err)
			continue
		}
		solrStatus := s.getLiveNodesStatuses(liveNodes, clusterState)

		problems := solrcheck.FindClusterProblems(s.ZooClient, clusterState, liveNodes)
		problems = append(problems, solrcheck.FindCloudStatusProblems(solrStatus, liveNodes)...)
		if len(problems) > 0 {
			for _, p := range problems {
				s.Logger.Infof("PROBLEM: %v", p)
			}
			if clusterStateGolden {
				s.AlertLog.Errorf("cluster state became not golden; see logs for details")
				clusterStateGolden = false
			}
			s.setStatusOp("cluster state is not golden; waiting for cluster state to become golden")
			s.Logger.Warningf("cluster state is not golden, skipping; see logs for details")
			continue
		}

		s.clearStatusOp()

		if !clusterStateGolden {
			s.AlertLog.Infof("cluster state became golden; resuming operation")
			clusterStateGolden = true
		}

		badlyBalancedOrgs := flagBadlyBalancedOrgs(s.Logger, s, solrStatus)
		if len(badlyBalancedOrgs) > 0 {
			s.Logger.Warningf(fmt.Sprintf("There are %d orgs with badly balanced shards.", len(badlyBalancedOrgs)))
		}

		if s.Storage.IsSplitsDisabled() {
			s.Logger.Infof("solrman splits are disabled")
		} else {
			shardSplits := computeShardSplits(s, solrStatus)
			anySplits := false
			for _, shardSplit := range shardSplits {
				if _, ok := badlyBalancedOrgs[shardSplit.Collection]; ok {
					s.Logger.Warningf("skipping split for badly balanced org %s_%s", shardSplit.Collection, shardSplit.Shard)
					continue
				}

				s.Logger.Infof("Scheduling split operation %v", shardSplit)
				result, err := s.SplitShard(shardSplit)
				if err != nil {
					s.Logger.Errorf("failed to schedule autogenerated split %+v: %s", shardSplit, err)
				} else if result.Error != "" {
					s.Logger.Warningf("failed to schedule autogenerated split %+v: %s", shardSplit, result.Error)
				} else {
					s.Logger.Infof("scheduled autogenerated split %+v", shardSplit)
					anySplits = true
				}
			}

			if anySplits {
				continue
			}
		}

		if s.Storage.IsMovesDisabled() {
			s.Logger.Infof("solrman moves are disabled")
		} else {
			// Long-running computation!
			s.setStatusOp("computing shard moves")
			shardMoves, err := computeShardMoves(solrStatus, evacuatingNodes, movesPerCycle)
			if err != nil {
				s.setStatusOp("failed to compute shard moves")
				s.Logger.Errorf("failed to compute shard moves: %s", err)
				continue
			}

			if len(shardMoves) == 0 {
				// Sleep an extra long time if everything is gucci.
				s.setStatusOp("nothing to do, everything is well balanced!")
				s.Logger.Debugf("nothing to do, everything is well balanced!")
				time.Sleep(quiescenceSleep - iterationSleep)
				continue
			}

			// Double-check we're in a good state before queuing the moves.
			if s.ZooClient.State() != zk.StateHasSession {
				s.setStatusOp("not connected to zk")
				s.Logger.Warningf("not connected to zk")
				continue
			}

			if s.hasInProgressOps() {
				s.clearStatusOp()
				s.Logger.Debugf("in progress ops")
				continue
			}

			for _, shardMove := range shardMoves {
				result, err := s.MoveShard(shardMove)
				if err != nil {
					s.Logger.Errorf("failed to schedule autogenerated move %+v: %s", shardMove, err)
				} else if result.Error != "" {
					s.Logger.Warningf("failed to schedule autogenerated move %+v: %s", shardMove, result.Error)
				} else {
					s.Logger.Infof("scheduled autogenerated move %+v", shardMove)
				}
			}
		}
	}
}

func (s *SolrManService) hasInProgressOps() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.inProgressOps) > 0
}

// For admin visibility
func (s *SolrManService) clearStatusOp() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.statusOp = nil
}

// For admin visibility
func (s *SolrManService) setStatusOp(status string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.statusOp = &solrmanapi.OpRecord{
		StartedMs: nowMillis(),
		Operation: solrmanapi.OpStatus,
		Error:     status,
	}
}

type SplitShardRequestWithSize struct {
	solrmanapi.SplitShardRequest
	NumDocs int64
}

type byNumDocsDesc []*SplitShardRequestWithSize

func (s byNumDocsDesc) Len() int {
	return len(s)
}
func (s byNumDocsDesc) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s byNumDocsDesc) Less(i, j int) bool {
	return s[i].NumDocs > s[j].NumDocs
}

// MaxInt64 returns the maximum of int64 values.
func MaxInt64(a, b int64) int64 {
	if a > b {
		return a
	}
	return b
}

// MinInt64 returns the minimum of int64 values.
func MinInt64(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}

func flagBadlyBalancedOrgs(logger smutil.Logger, s *SolrManService, clusterState solrmanapi.SolrCloudStatus) map[string]bool {
	orgToMin := make(map[string]int64)
	orgToMax := make(map[string]int64)

	getOrg := func(coreName string) string {
		return strings.Split(coreName, "_")[0]
	}

	for _, v := range clusterState {
		for coreName, status := range v.Cores {
			org := getOrg(coreName)

			if min, ok := orgToMin[org]; ok {
				orgToMin[org] = MinInt64(min, status.NumDocs)
			} else {
				orgToMin[org] = status.NumDocs
			}

			if max, ok := orgToMax[org]; ok {
				orgToMax[org] = MaxInt64(max, status.NumDocs)
			} else {
				orgToMax[org] = status.NumDocs
			}
		}
	}

	badlyBalancedOrgs := make(map[string]bool)
	for org, max := range orgToMax {
		if (2+float64(orgToMin[org]))/(2+float64(max)) < allowedMinToMaxShardSizeRatio {
			badlyBalancedOrgs[org] = true
			logger.Warningf("Shards are getting imbalanced for org: " + org)
		}
	}

	return badlyBalancedOrgs
}

func computeShardSplits(s *SolrManService, clusterState solrmanapi.SolrCloudStatus) []*solrmanapi.SplitShardRequest {
	machineToSplitOps := make(map[string][]*SplitShardRequestWithSize)
	anyShardTooBig := false
	for machine, v := range clusterState {
		for _, status := range v.Cores {
			// Continue if this collection has too many shards i.e greater than maxShardsPerMachine * num solr machines
			if collstate, err := s.SolrMonitor.GetCollectionState(status.Collection); err != nil {
				continue
			} else if len(collstate.Shards) > maxShardsPerMachine*len(clusterState) {
				continue
			}

			if status.NumDocs > allSplitsDocCountTrigger {
				anyShardTooBig = true
			}

			if status.NumDocs > splitShardsWithDocCount {
				x := &SplitShardRequestWithSize{
					solrmanapi.SplitShardRequest{
						Collection: status.Collection,
						Shard:      status.Shard,
					},
					status.NumDocs,
				}
				machineToSplitOps[machine] = append(machineToSplitOps[machine], x)
			}
		}
	}

	if !anyShardTooBig {
		return nil
	}

	// sort by biggest to smallest shards
	// keep biggest splitsPerMachine shards only
	// flatten and return

	splitOps := make([]*solrmanapi.SplitShardRequest, 0)
	for _, ops := range machineToSplitOps {
		sort.Sort(byNumDocsDesc(ops))
		for i, op := range ops {
			if i == splitsPerMachine {
				break
			}
			splitOps = append(splitOps, &op.SplitShardRequest)
		}
	}

	return splitOps
}

func computeShardMoves(clusterState solrmanapi.SolrCloudStatus, evacuatingNodes []string, count int) ([]*solrmanapi.MoveShardRequest, error) {
	baseModel, err := createModel(clusterState, evacuatingNodes)
	if err != nil {
		return nil, err
	}

	// Leave 1 cores open for handling queries / etc.
	numCPU := runtime.GOMAXPROCS(0) - 1
	if numCPU < 1 {
		numCPU = 1
	}

	var moves []*smmodel.Move
	immobileCores := map[string]bool{}
	m := baseModel
	for i := 0; i < count; i++ {
		mPrime, move := m.ComputeNextMove(numCPU, immobileCores)
		if m == mPrime {
			break
		}
		moves = append(moves, move)
		immobileCores[move.Core.Name] = true
		m = mPrime
	}

	var shardMoves []*solrmanapi.MoveShardRequest
	if len(moves) > 0 {
		shardMoves = make([]*solrmanapi.MoveShardRequest, len(moves))
		for i, m := range moves {
			shardMoves[i] = &solrmanapi.MoveShardRequest{
				Collection: m.Core.Collection,
				Shard:      m.Core.Shard,
				SrcNode:    m.FromNode.Address,
				DstNode:    m.ToNode.Address,
			}
		}
	}

	return shardMoves, nil
}

func createModel(clusterState solrmanapi.SolrCloudStatus, evacuatingNodes []string) (*smmodel.Model, error) {
	var currentNode *smmodel.Node
	m := &smmodel.Model{}
	seenNodeNames := map[string]bool{}
	collectionMap := make(map[string]*smmodel.Collection)
	evacuatingNodeSet := map[string]bool{}
	for _, n := range evacuatingNodes {
		evacuatingNodeSet[n] = true
	}

	for _, nodeStatus := range clusterState {
		if seenNodeNames[nodeStatus.NodeName] {
			return nil, smutil.Errorf("already seen: %v", nodeStatus.NodeName)
		}
		seenNodeNames[nodeStatus.NodeName] = true
		currentNode = &smmodel.Node{
			Name:       nodeStatus.Hostname,
			Address:    nodeStatus.NodeName,
			Evacuating: evacuatingNodeSet[nodeStatus.Hostname],
		}
		m.AddNode(currentNode)

		for _, coreStatus := range nodeStatus.Cores {
			collName := coreStatus.Collection
			collection := collectionMap[collName]
			if collection == nil {
				collection = &smmodel.Collection{Name: collName}
				collectionMap[collName] = collection
				m.Collections = append(m.Collections, collection)
			}

			core := &smmodel.Core{
				Name:       coreStatus.Name,
				Collection: collName,
				Shard:      coreStatus.Shard,
				Docs:       float64(coreStatus.NumDocs),
				Size:       float64(coreStatus.IndexSize),
			}
			collection.Add(core)
			currentNode.Add(core)
			m.Add(core)
		}
	}
	return m, nil
}
<commit_msg>Recheck solrman disable after computing moves<commit_after>// Copyright 2016 FullStory, Inc.
//
// 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 smservice

import (
	"fmt"
	"runtime"
	"sort"
	"strings"
	"time"

	"github.com/fullstorydev/gosolr/smutil"
	"github.com/fullstorydev/gosolr/solrcheck"
	"github.com/fullstorydev/gosolr/solrman/smmodel"
	"github.com/fullstorydev/gosolr/solrman/solrmanapi"
	"github.com/samuel/go-zookeeper/zk"
)

const (
	movesPerCycle                 = 10               // How many shards to move at a time
	iterationSleep                = 1 * time.Minute  // How long to sleep each attempt, no matter what
	quiescenceSleep               = 10 * time.Minute // How long to sleep when stability is reached
	splitsPerMachine              = 4                // How many shard splitting jobs to schedule on 1 physical machine at a time
	splitShardsWithDocCount       = 4000000          // Split shards with doc count > this
	allSplitsDocCountTrigger      = 4004000          // But don't do any splits until at least one shard > this
	allowedMinToMaxShardSizeRatio = 0.2              // Ratio of smallest shard to biggest shard < this then warn about imbalance
	maxShardsPerMachine           = 16               // Maximum number of shards per machine in the cluster
)

// Runs the main solr management loop, never returns.
func (s *SolrManService) RunSolrMan() {
	s.setStatusOp("solrman is starting up")
	clusterStateGolden := true // assume true to start

	first := true
	for {
		if !first {
			time.Sleep(iterationSleep)
		}
		first = false

		if s.ZooClient.State() != zk.StateHasSession {
			s.setStatusOp("not connected to zk")
			s.Logger.Warningf("not connected to zk")
			continue
		}

		if s.hasInProgressOps() {
			s.clearStatusOp()
			s.Logger.Debugf("in progress ops")
			continue
		}

		if s.Storage.IsDisabled() {
			s.setStatusOp("solrman is disabled")
			s.Logger.Infof("solrman is disabled")
			continue
		}

		evacuatingNodes, err := s.Storage.GetEvacuateNodeList()
		if err != nil {
			s.setStatusOp("failed to determine hosts to evacuate")
			s.Logger.Errorf("failed to determine hosts to evacuate: %s", err)
			continue
		}

		clusterState, err := s.SolrMonitor.GetCurrentState()
		if err != nil {
			s.setStatusOp("failed to retrieve cluster state")
			s.Logger.Errorf("failed to retrieve cluster state: %s", err)
			continue
		}
		liveNodes, err := s.SolrMonitor.GetLiveNodes()
		if err != nil {
			s.setStatusOp("failed to retrieve live nodes")
			s.Logger.Errorf("failed to retrieve live nodes: %s", err)
			continue
		}
		solrStatus := s.getLiveNodesStatuses(liveNodes, clusterState)

		problems := solrcheck.FindClusterProblems(s.ZooClient, clusterState, liveNodes)
		problems = append(problems, solrcheck.FindCloudStatusProblems(solrStatus, liveNodes)...)
		if len(problems) > 0 {
			for _, p := range problems {
				s.Logger.Infof("PROBLEM: %v", p)
			}
			if clusterStateGolden {
				s.AlertLog.Errorf("cluster state became not golden; see logs for details")
				clusterStateGolden = false
			}
			s.setStatusOp("cluster state is not golden; waiting for cluster state to become golden")
			s.Logger.Warningf("cluster state is not golden, skipping; see logs for details")
			continue
		}

		s.clearStatusOp()

		if !clusterStateGolden {
			s.AlertLog.Infof("cluster state became golden; resuming operation")
			clusterStateGolden = true
		}

		badlyBalancedOrgs := flagBadlyBalancedOrgs(s.Logger, s, solrStatus)
		if len(badlyBalancedOrgs) > 0 {
			s.Logger.Warningf(fmt.Sprintf("There are %d orgs with badly balanced shards.", len(badlyBalancedOrgs)))
		}

		if s.Storage.IsSplitsDisabled() {
			s.Logger.Infof("solrman splits are disabled")
		} else {
			shardSplits := computeShardSplits(s, solrStatus)
			anySplits := false
			for _, shardSplit := range shardSplits {
				if _, ok := badlyBalancedOrgs[shardSplit.Collection]; ok {
					s.Logger.Warningf("skipping split for badly balanced org %s_%s", shardSplit.Collection, shardSplit.Shard)
					continue
				}

				s.Logger.Infof("Scheduling split operation %v", shardSplit)
				result, err := s.SplitShard(shardSplit)
				if err != nil {
					s.Logger.Errorf("failed to schedule autogenerated split %+v: %s", shardSplit, err)
				} else if result.Error != "" {
					s.Logger.Warningf("failed to schedule autogenerated split %+v: %s", shardSplit, result.Error)
				} else {
					s.Logger.Infof("scheduled autogenerated split %+v", shardSplit)
					anySplits = true
				}
			}

			if anySplits {
				continue
			}
		}

		if s.Storage.IsMovesDisabled() {
			s.Logger.Infof("solrman moves are disabled")
		} else {
			// Long-running computation!
			s.setStatusOp("computing shard moves")
			shardMoves, err := computeShardMoves(solrStatus, evacuatingNodes, movesPerCycle)
			if err != nil {
				s.setStatusOp("failed to compute shard moves")
				s.Logger.Errorf("failed to compute shard moves: %s", err)
				continue
			}

			if len(shardMoves) == 0 {
				// Sleep an extra long time if everything is gucci.
				s.setStatusOp("nothing to do, everything is well balanced!")
				s.Logger.Debugf("nothing to do, everything is well balanced!")
				time.Sleep(quiescenceSleep - iterationSleep)
				continue
			}

			// Computing moves takes a long time; double-check we're in a good state before starting the moves.
			if s.ZooClient.State() != zk.StateHasSession {
				s.setStatusOp("not connected to zk")
				s.Logger.Warningf("not connected to zk")
				continue
			}

			if s.hasInProgressOps() {
				s.clearStatusOp()
				s.Logger.Debugf("in progress ops")
				continue
			}

			if s.Storage.IsDisabled() {
				s.setStatusOp("solrman is disabled")
				s.Logger.Infof("solrman is disabled")
				continue
			}

			if s.Storage.IsMovesDisabled() {
				s.clearStatusOp()
				s.Logger.Infof("solrman moves are disabled")
				continue
			}

			for _, shardMove := range shardMoves {
				result, err := s.MoveShard(shardMove)
				if err != nil {
					s.Logger.Errorf("failed to schedule autogenerated move %+v: %s", shardMove, err)
				} else if result.Error != "" {
					s.Logger.Warningf("failed to schedule autogenerated move %+v: %s", shardMove, result.Error)
				} else {
					s.Logger.Infof("scheduled autogenerated move %+v", shardMove)
				}
			}
		}
	}
}

func (s *SolrManService) hasInProgressOps() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.inProgressOps) > 0
}

// For admin visibility
func (s *SolrManService) clearStatusOp() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.statusOp = nil
}

// For admin visibility
func (s *SolrManService) setStatusOp(status string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.statusOp = &solrmanapi.OpRecord{
		StartedMs: nowMillis(),
		Operation: solrmanapi.OpStatus,
		Error:     status,
	}
}

type SplitShardRequestWithSize struct {
	solrmanapi.SplitShardRequest
	NumDocs int64
}

type byNumDocsDesc []*SplitShardRequestWithSize

func (s byNumDocsDesc) Len() int {
	return len(s)
}
func (s byNumDocsDesc) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s byNumDocsDesc) Less(i, j int) bool {
	return s[i].NumDocs > s[j].NumDocs
}

// MaxInt64 returns the maximum of int64 values.
func MaxInt64(a, b int64) int64 {
	if a > b {
		return a
	}
	return b
}

// MinInt64 returns the minimum of int64 values.
func MinInt64(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}

func flagBadlyBalancedOrgs(logger smutil.Logger, s *SolrManService, clusterState solrmanapi.SolrCloudStatus) map[string]bool {
	orgToMin := make(map[string]int64)
	orgToMax := make(map[string]int64)

	getOrg := func(coreName string) string {
		return strings.Split(coreName, "_")[0]
	}

	for _, v := range clusterState {
		for coreName, status := range v.Cores {
			org := getOrg(coreName)

			if min, ok := orgToMin[org]; ok {
				orgToMin[org] = MinInt64(min, status.NumDocs)
			} else {
				orgToMin[org] = status.NumDocs
			}

			if max, ok := orgToMax[org]; ok {
				orgToMax[org] = MaxInt64(max, status.NumDocs)
			} else {
				orgToMax[org] = status.NumDocs
			}
		}
	}

	badlyBalancedOrgs := make(map[string]bool)
	for org, max := range orgToMax {
		if (2+float64(orgToMin[org]))/(2+float64(max)) < allowedMinToMaxShardSizeRatio {
			badlyBalancedOrgs[org] = true
			logger.Warningf("Shards are getting imbalanced for org: " + org)
		}
	}

	return badlyBalancedOrgs
}

func computeShardSplits(s *SolrManService, clusterState solrmanapi.SolrCloudStatus) []*solrmanapi.SplitShardRequest {
	machineToSplitOps := make(map[string][]*SplitShardRequestWithSize)
	anyShardTooBig := false
	for machine, v := range clusterState {
		for _, status := range v.Cores {
			// Continue if this collection has too many shards i.e greater than maxShardsPerMachine * num solr machines
			if collstate, err := s.SolrMonitor.GetCollectionState(status.Collection); err != nil {
				continue
			} else if len(collstate.Shards) > maxShardsPerMachine*len(clusterState) {
				continue
			}

			if status.NumDocs > allSplitsDocCountTrigger {
				anyShardTooBig = true
			}

			if status.NumDocs > splitShardsWithDocCount {
				x := &SplitShardRequestWithSize{
					solrmanapi.SplitShardRequest{
						Collection: status.Collection,
						Shard:      status.Shard,
					},
					status.NumDocs,
				}
				machineToSplitOps[machine] = append(machineToSplitOps[machine], x)
			}
		}
	}

	if !anyShardTooBig {
		return nil
	}

	// sort by biggest to smallest shards
	// keep biggest splitsPerMachine shards only
	// flatten and return

	splitOps := make([]*solrmanapi.SplitShardRequest, 0)
	for _, ops := range machineToSplitOps {
		sort.Sort(byNumDocsDesc(ops))
		for i, op := range ops {
			if i == splitsPerMachine {
				break
			}
			splitOps = append(splitOps, &op.SplitShardRequest)
		}
	}

	return splitOps
}

func computeShardMoves(clusterState solrmanapi.SolrCloudStatus, evacuatingNodes []string, count int) ([]*solrmanapi.MoveShardRequest, error) {
	baseModel, err := createModel(clusterState, evacuatingNodes)
	if err != nil {
		return nil, err
	}

	// Leave 1 cores open for handling queries / etc.
	numCPU := runtime.GOMAXPROCS(0) - 1
	if numCPU < 1 {
		numCPU = 1
	}

	var moves []*smmodel.Move
	immobileCores := map[string]bool{}
	m := baseModel
	for i := 0; i < count; i++ {
		mPrime, move := m.ComputeNextMove(numCPU, immobileCores)
		if m == mPrime {
			break
		}
		moves = append(moves, move)
		immobileCores[move.Core.Name] = true
		m = mPrime
	}

	var shardMoves []*solrmanapi.MoveShardRequest
	if len(moves) > 0 {
		shardMoves = make([]*solrmanapi.MoveShardRequest, len(moves))
		for i, m := range moves {
			shardMoves[i] = &solrmanapi.MoveShardRequest{
				Collection: m.Core.Collection,
				Shard:      m.Core.Shard,
				SrcNode:    m.FromNode.Address,
				DstNode:    m.ToNode.Address,
			}
		}
	}

	return shardMoves, nil
}

func createModel(clusterState solrmanapi.SolrCloudStatus, evacuatingNodes []string) (*smmodel.Model, error) {
	var currentNode *smmodel.Node
	m := &smmodel.Model{}
	seenNodeNames := map[string]bool{}
	collectionMap := make(map[string]*smmodel.Collection)
	evacuatingNodeSet := map[string]bool{}
	for _, n := range evacuatingNodes {
		evacuatingNodeSet[n] = true
	}

	for _, nodeStatus := range clusterState {
		if seenNodeNames[nodeStatus.NodeName] {
			return nil, smutil.Errorf("already seen: %v", nodeStatus.NodeName)
		}
		seenNodeNames[nodeStatus.NodeName] = true
		currentNode = &smmodel.Node{
			Name:       nodeStatus.Hostname,
			Address:    nodeStatus.NodeName,
			Evacuating: evacuatingNodeSet[nodeStatus.Hostname],
		}
		m.AddNode(currentNode)

		for _, coreStatus := range nodeStatus.Cores {
			collName := coreStatus.Collection
			collection := collectionMap[collName]
			if collection == nil {
				collection = &smmodel.Collection{Name: collName}
				collectionMap[collName] = collection
				m.Collections = append(m.Collections, collection)
			}

			core := &smmodel.Core{
				Name:       coreStatus.Name,
				Collection: collName,
				Shard:      coreStatus.Shard,
				Docs:       float64(coreStatus.NumDocs),
				Size:       float64(coreStatus.IndexSize),
			}
			collection.Add(core)
			currentNode.Add(core)
			m.Add(core)
		}
	}
	return m, nil
}
<|endoftext|>


<commit_before>package aci

import (
	"testing"
)

func TestVlanpool1(t *testing.T) {
	vlanspoolSplitTest(t, "", "", "")
	vlanspoolSplitTest(t, "vlanns-[a]-b", "a", "b")
	vlanspoolSplitTest(t, "vlanns-[]-", "", "")
	vlanspoolSplitTest(t, "vlanns-[a]-", "a", "")
	vlanspoolSplitTest(t, "vlanns-[]-b", "", "b")
	vlanspoolSplitTest(t, "vlanns-[a-b", "a", "b")
	vlanspoolSplitTest(t, "vlanns-a]-b", "a", "b")
	vlanspoolSplitTest(t, "vlanns-[-b", "", "b")
	vlanspoolSplitTest(t, "vlanns-]-b", "", "b")
}

func vlanspoolSplitTest(t *testing.T, input, wantPool, wantMode string) {
	resultPool, resultMode := vlanpoolSplit(input)
	if resultPool != wantPool || resultMode != wantMode {
		t.Errorf("input=%s wantPool=%s gotPool=%s wantMode=%s gotMode=%s", input, wantPool, resultPool, wantMode, resultMode)
	}
}
<commit_msg>Tests.<commit_after>package aci

import (
	"testing"
)

func TestVlanpool1(t *testing.T) {
	vlanspoolSplitTest(t, "", "", "")
	vlanspoolSplitTest(t, "vlanns-[a]-b", "a", "b")
	vlanspoolSplitTest(t, "vlanns-[]-", "", "")
	vlanspoolSplitTest(t, "vlanns-[a]-", "a", "")
	vlanspoolSplitTest(t, "vlanns-[]-b", "", "b")
	vlanspoolSplitTest(t, "vlanns-[a-b", "a", "b")
	vlanspoolSplitTest(t, "vlanns-a]-b", "a", "b")
	vlanspoolSplitTest(t, "vlanns-[-b", "", "b")
	vlanspoolSplitTest(t, "vlanns-]-b", "", "b")
	vlanspoolSplitTest(t, "vlanns-a-b", "a", "b")
	vlanspoolSplitTest(t, "vlanns--", "", "")
}

func vlanspoolSplitTest(t *testing.T, input, wantPool, wantMode string) {
	resultPool, resultMode := vlanpoolSplit(input)
	if resultPool != wantPool || resultMode != wantMode {
		t.Errorf("input=%s wantPool=%s gotPool=%s wantMode=%s gotMode=%s", input, wantPool, resultPool, wantMode, resultMode)
	}
}
<|endoftext|>


<commit_before>package pool

import (
	"fmt"
	strfmt "github.com/go-openapi/strfmt"
	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"

	bslcommon "github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/common"

	sl "github.com/maximilien/softlayer-go/softlayer"
	"github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/pool/client"
	operations "github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/pool/client/vm"

	. "github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/vm"

	"github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/pool/models"
)

const SOFTLAYER_POOL_DELETER_LOG_TAG = "SoftLayerPoolDeleter"

type softLayerPoolDeleter struct {
	softLayerClient        sl.Client
	softLayerVmPoolClient  *client.SoftLayerVMPool
	logger       boshlog.Logger
}

func NewSoftLayerPoolDeleter(softLayerVmPoolClient *client.SoftLayerVMPool, softLayerClient sl.Client, logger boshlog.Logger) VMDeleter {
	return &softLayerPoolDeleter{
		softLayerClient:        softLayerClient,
		softLayerVmPoolClient: softLayerVmPoolClient,
		logger: logger,
	}
}

func (c *softLayerPoolDeleter) Delete(cid int) error {
	getVmByCidResp, err := c.softLayerVmPoolClient.VM.GetVMByCid(operations.NewGetVMByCidParams().WithCid(int32(cid)))
	if getVmByCidResp.Payload.VM != nil {
		free := models.VMState{
			State: models.StateFree,
		}
		_, err = c.softLayerVmPoolClient.VM.UpdateVMWithState(operations.NewUpdateVMWithStateParams().WithBody(&free).WithCid(int32(cid)))
		if err != nil {
			return bosherr.WrapErrorf(err, "Updating state of vm %d in pool to free", cid)
		}
		return nil
	}

	virtualGuest, err := bslcommon.GetObjectDetailsOnVirtualGuest(c.softLayerClient, cid)
	if err != nil {
		return bosherr.WrapError(err, fmt.Sprintf("Getting virtual guest %d details from SoftLayer", cid))
	}

	slPoolVm := &models.VM{
		Cid: int32(cid),
		CPU: int32(virtualGuest.StartCpus),
		MemoryMb: int32(virtualGuest.MaxMemory),
		IP:  strfmt.IPv4(virtualGuest.PrimaryBackendIpAddress),
		Hostname: virtualGuest.FullyQualifiedDomainName,
		PrivateVlan: int32(virtualGuest.PrimaryBackendNetworkComponent.NetworkVlan.Id),
		PublicVlan: int32(virtualGuest.PrimaryNetworkComponent.NetworkVlan.Id),
		State: models.StateFree,
	}
	_, err = c.softLayerVmPoolClient.VM.AddVM(operations.NewAddVMParams().WithBody(slPoolVm))
	if err != nil {
		return bosherr.WrapError(err, fmt.Sprintf("Adding vm %d to pool", cid))
	}

	return nil
}<commit_msg>updated softlayer_pool_deleter according to the change of vps server<commit_after>package pool

import (
	"fmt"
	strfmt "github.com/go-openapi/strfmt"
	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"

	bslcommon "github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/common"

	sl "github.com/maximilien/softlayer-go/softlayer"
	"github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/pool/client"
	operations "github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/pool/client/vm"

	. "github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/vm"

	"github.com/cloudfoundry/bosh-softlayer-cpi/softlayer/pool/models"
)

const SOFTLAYER_POOL_DELETER_LOG_TAG = "SoftLayerPoolDeleter"

type softLayerPoolDeleter struct {
	softLayerClient        sl.Client
	softLayerVmPoolClient  *client.SoftLayerVMPool
	logger       boshlog.Logger
}

func NewSoftLayerPoolDeleter(softLayerVmPoolClient *client.SoftLayerVMPool, softLayerClient sl.Client, logger boshlog.Logger) VMDeleter {
	return &softLayerPoolDeleter{
		softLayerClient:        softLayerClient,
		softLayerVmPoolClient: softLayerVmPoolClient,
		logger: logger,
	}
}

func (c *softLayerPoolDeleter) Delete(cid int) error {
	_, err := c.softLayerVmPoolClient.VM.GetVMByCid(operations.NewGetVMByCidParams().WithCid(int32(cid)))
	if err != nil {
		_, ok := err.(*operations.DeleteVMNotFound)
		if ok {
			virtualGuest, err := bslcommon.GetObjectDetailsOnVirtualGuest(c.softLayerClient, cid)
			if err != nil {
				return bosherr.WrapError(err, fmt.Sprintf("Getting virtual guest %d details from SoftLayer", cid))
			}

			slPoolVm := &models.VM{
				Cid: int32(cid),
				CPU: int32(virtualGuest.StartCpus),
				MemoryMb: int32(virtualGuest.MaxMemory),
				IP:  strfmt.IPv4(virtualGuest.PrimaryBackendIpAddress),
				Hostname: virtualGuest.FullyQualifiedDomainName,
				PrivateVlan: int32(virtualGuest.PrimaryBackendNetworkComponent.NetworkVlan.Id),
				PublicVlan: int32(virtualGuest.PrimaryNetworkComponent.NetworkVlan.Id),
				State: models.StateFree,
			}
			_, err = c.softLayerVmPoolClient.VM.AddVM(operations.NewAddVMParams().WithBody(slPoolVm))
			if err != nil {
				return bosherr.WrapError(err, fmt.Sprintf("Adding vm %d to pool", cid))
			}
			return nil
		}
		return bosherr.WrapError(err, "Removing vm from pool")
	}

	free := models.VMState{
		State: models.StateFree,
	}
	_, err = c.softLayerVmPoolClient.VM.UpdateVMWithState(operations.NewUpdateVMWithStateParams().WithBody(&free).WithCid(int32(cid)))
	if err != nil {
		return bosherr.WrapErrorf(err, "Updating state of vm %d in pool to free", cid)
	}

	return nil
}<|endoftext|>


<commit_before>package modbus

import (
	"errors"
	"os"
	"time"
)

func ModbusRead(fd *os.File, addr byte, sr, nr uint16, data []byte) ([]byte, error) {
	//Preparation for Sending a Packet
	var send_packet = make([]byte, 8)

	//Packet Construction
	send_packet[0] = addr            // Slave Address
	send_packet[1] = 0x03            // Function Code 0x03 = Multiple Read
	send_packet[2] = byte(sr >> 8)   // Start Register (High Byte)
	send_packet[3] = byte(sr & 0xff) // Start Register (Low Byte)
	send_packet[4] = byte(nr >> 8)   // Number of Registers (High Byte)
	send_packet[5] = byte(nr & 0xff) // Number of Registers (Low Byte)

	//Add CRC16
	send_packet_crc := ModbusCrc(send_packet[:6])
	send_packet[6] = byte(send_packet_crc & 0xff)
	send_packet[7] = byte(send_packet_crc >> 8)

	// Preparation for Receiving a Packet
	var recv_packet = make([]byte, 256)
	_, err := fd.Write(send_packet)
	if err != nil {
		return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION")
	}
	time.Sleep(300 * time.Millisecond)
	_, err = fd.Read(recv_packet)
	if err != nil {
		return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION")
	}

	// Parse the Response
	if recv_packet[0] != send_packet[0] || recv_packet[1] != send_packet[1] {
		if recv_packet[0] == send_packet[0] && recv_packet[1]&0x7f == send_packet[1] {
			switch recv_packet[2] {
			case 1:
				return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_FUNCTION")
			case 2:
				return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_ADDRESS")
			case 3:
				return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_VALUE")
			case 4:
				return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_OPERATION")
			}
		}
		return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION")
	}

	//CRC check
	l := recv_packet[2]
	recv_packet_crc := ModbusCrc(recv_packet[:3+l])
	if recv_packet[3+l] != byte((recv_packet_crc&0xff)) || recv_packet[3+l+1] != byte((recv_packet_crc>>8)) {
		return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION")
	}
	return recv_packet[3 : l+3], nil
}

func ModbusWrite(fd *os.File, addr byte, sr, nr uint16, data []byte) error {
	var send_packet = make([]byte, 256)

	// Packet Construction
	send_packet[0] = addr            // Slave Address
	send_packet[1] = 0x10            // Function Code 0x10 = Multiple Write
	send_packet[2] = byte(sr >> 8)   // Start Register (High Byte)
	send_packet[3] = byte(sr & 0xff) // Start Register (Low Byte)
	send_packet[4] = byte(nr >> 8)   // Number of Registers (High Byte)
	send_packet[5] = byte(nr & 0xff) // Number of Registers (Low Byte)
	send_packet[6] = byte(nr * 2)

	for i := 0; i < int((nr * 2)); i++ {
		send_packet[7+i] = data[i]
	}

	length := 7 + nr*2 + 2
	// Add CRC16
	send_packet_crc := ModbusCrc(send_packet[:length-2])
	send_packet[length-2] = byte(send_packet_crc & 0xff)
	send_packet[length-1] = byte(send_packet_crc >> 8)

	// Preparation for Receiving a Packet
	var recv_packet = make([]byte, 256)
	_, err := fd.Write(send_packet)
	if err != nil {
		return errors.New("MODBUS_ERROR_COMMUNICATION")
	}
	time.Sleep(300 * time.Millisecond)
	_, err = fd.Read(recv_packet)
	if err != nil {
		return errors.New("MODBUS_ERROR_COMMUNICATION")
	}

	// Parse the Response
	if recv_packet[0] != send_packet[0] || recv_packet[1] != send_packet[1] {
		if recv_packet[0] == send_packet[0] && recv_packet[1]&0x7f == send_packet[1] {
			switch recv_packet[2] {
			case 1:
				return errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_FUNCTION")
			case 2:
				return errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_ADDRESS")
			case 3:
				return errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_VALUE")
			case 4:
				return errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_OPERATION")
			}
		}
		return errors.New("MODBUS_ERROR_COMMUNICATION")
	}

	//Target Data Filed Check
	if recv_packet[2] == send_packet[2] && recv_packet[3] == send_packet[3] && recv_packet[4] == send_packet[4] && recv_packet[5] == send_packet[5] {
		//CRC check
		recv_packet_crc := ModbusCrc(recv_packet[:6])
		if recv_packet[6] == byte((recv_packet_crc&0xff)) && recv_packet[7] == byte((recv_packet_crc>>8)) {
			return nil
		}
	}
	return errors.New("MODBUS_ERROR_COMMUNICATION")
}

func ModbusCrc(data []byte) uint16 {
	var crc16 uint16 = 0xffff
	l := len(data)
	for i := 0; i < l; i++ {
		crc16 ^= uint16(data[i])
		for j := 0; j < 8; j++ {
			if crc16&0x0001 > 0 {
				crc16 = (crc16 >> 1) ^ 0xA001
			} else {
				crc16 >>= 1
			}
		}
	}
	return crc16
}
<commit_msg>add annotion<commit_after>package modbus

import (
	"errors"
	"os"
	"time"
)
/**
 *  Modbus Read
 *
 *  * parameters
 *   int        fd:  file descripter for serial device
 *   uint8_t  addr:  slave device address
 *   uint16_t sr:    starting register number
 *   uint16_t nr:    number of registers to read or write
 *   uint8_t data[]: memory area for read data -- the size of footprint must be nr*2
 */
func ModbusRead(fd *os.File, addr byte, sr, nr uint16, data []byte) ([]byte, error) {
	//Preparation for Sending a Packet
	var send_packet = make([]byte, 8)

	//Packet Construction
	send_packet[0] = addr            // Slave Address
	send_packet[1] = 0x03            // Function Code 0x03 = Multiple Read
	send_packet[2] = byte(sr >> 8)   // Start Register (High Byte)
	send_packet[3] = byte(sr & 0xff) // Start Register (Low Byte)
	send_packet[4] = byte(nr >> 8)   // Number of Registers (High Byte)
	send_packet[5] = byte(nr & 0xff) // Number of Registers (Low Byte)

	//Add CRC16
	send_packet_crc := ModbusCrc(send_packet[:6])
	send_packet[6] = byte(send_packet_crc & 0xff)
	send_packet[7] = byte(send_packet_crc >> 8)

	// Preparation for Receiving a Packet
	var recv_packet = make([]byte, 256)
	_, err := fd.Write(send_packet)
	if err != nil {
		return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION")
	}
	time.Sleep(300 * time.Millisecond)
	_, err = fd.Read(recv_packet)
	if err != nil {
		return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION")
	}

	// Parse the Response
	if recv_packet[0] != send_packet[0] || recv_packet[1] != send_packet[1] {
		if recv_packet[0] == send_packet[0] && recv_packet[1]&0x7f == send_packet[1] {
			switch recv_packet[2] {
			case 1:
				return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_FUNCTION")
			case 2:
				return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_ADDRESS")
			case 3:
				return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_VALUE")
			case 4:
				return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_OPERATION")
			}
		}
		return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION")
	}

	//CRC check
	l := recv_packet[2]
	recv_packet_crc := ModbusCrc(recv_packet[:3+l])
	if recv_packet[3+l] != byte((recv_packet_crc&0xff)) || recv_packet[3+l+1] != byte((recv_packet_crc>>8)) {
		return []byte{}, errors.New("MODBUS_ERROR_COMMUNICATION")
	}
	return recv_packet[3 : l+3], nil
}

/**
 *  Modbus Write
 *
 *  * parameters
 *   int        fd:  file descripter for serial device
 *   uint8_t  addr:  slave device address
 *   uint16_t sr:    starting register number
 *   uint16_t nr:    number of registers to write
 *   uint8_t data[]: memory area for writing data -- the size of footprint must be nr*2
 */
func ModbusWrite(fd *os.File, addr byte, sr, nr uint16, data []byte) error {
	var send_packet = make([]byte, 256)

	// Packet Construction
	send_packet[0] = addr            // Slave Address
	send_packet[1] = 0x10            // Function Code 0x10 = Multiple Write
	send_packet[2] = byte(sr >> 8)   // Start Register (High Byte)
	send_packet[3] = byte(sr & 0xff) // Start Register (Low Byte)
	send_packet[4] = byte(nr >> 8)   // Number of Registers (High Byte)
	send_packet[5] = byte(nr & 0xff) // Number of Registers (Low Byte)
	send_packet[6] = byte(nr * 2)

	for i := 0; i < int((nr * 2)); i++ {
		send_packet[7+i] = data[i]
	}

	length := 7 + nr*2 + 2
	// Add CRC16
	send_packet_crc := ModbusCrc(send_packet[:length-2])
	send_packet[length-2] = byte(send_packet_crc & 0xff)
	send_packet[length-1] = byte(send_packet_crc >> 8)

	// Preparation for Receiving a Packet
	var recv_packet = make([]byte, 256)
	_, err := fd.Write(send_packet)
	if err != nil {
		return errors.New("MODBUS_ERROR_COMMUNICATION")
	}
	time.Sleep(300 * time.Millisecond)
	_, err = fd.Read(recv_packet)
	if err != nil {
		return errors.New("MODBUS_ERROR_COMMUNICATION")
	}

	// Parse the Response
	if recv_packet[0] != send_packet[0] || recv_packet[1] != send_packet[1] {
		if recv_packet[0] == send_packet[0] && recv_packet[1]&0x7f == send_packet[1] {
			switch recv_packet[2] {
			case 1:
				return errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_FUNCTION")
			case 2:
				return errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_ADDRESS")
			case 3:
				return errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_VALUE")
			case 4:
				return errors.New("MODBUS_ERROR_COMMUNICATION_ILLEGAL_OPERATION")
			}
		}
		return errors.New("MODBUS_ERROR_COMMUNICATION")
	}

	//Target Data Filed Check
	if recv_packet[2] == send_packet[2] && recv_packet[3] == send_packet[3] && recv_packet[4] == send_packet[4] && recv_packet[5] == send_packet[5] {
		//CRC check
		recv_packet_crc := ModbusCrc(recv_packet[:6])
		if recv_packet[6] == byte((recv_packet_crc&0xff)) && recv_packet[7] == byte((recv_packet_crc>>8)) {
			return nil
		}
	}
	return errors.New("MODBUS_ERROR_COMMUNICATION")
}

func ModbusCrc(data []byte) uint16 {
	var crc16 uint16 = 0xffff
	l := len(data)
	for i := 0; i < l; i++ {
		crc16 ^= uint16(data[i])
		for j := 0; j < 8; j++ {
			if crc16&0x0001 > 0 {
				crc16 = (crc16 >> 1) ^ 0xA001
			} else {
				crc16 >>= 1
			}
		}
	}
	return crc16
}
<|endoftext|>


<commit_before>/*-
Copyright (c) 2016, Jörg Pernfuß <joerg.pernfuss@1und1.de>
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"time"

	"github.com/mjolnir42/scrypth64"
	"github.com/satori/go.uuid"
)

func (s *supervisor) activate_user(q *msg.Request) {
	result := msg.Result{Type: `supervisor`, Action: `activate_user`}
	var (
		timer                               *time.Timer
		plain                               []byte
		err                                 error
		kex                                 *auth.Kex
		validFrom, expiresAt, credExpiresAt time.Time
		token                               auth.Token
		userId                              string
		userUUID                            uuid.UUID
		ok                                  bool
		mcf                                 scrypth64.Mcf
		tx                                  *sql.Tx
	)
	data := q.Super.Data

	if s.readonly {
		result.Conflict(fmt.Errorf(`Readonly instance`))
		goto conflict
	}

	// start response timer
	timer = time.NewTimer(1 * time.Second)
	defer timer.Stop()

	// -> get kex
	if kex = s.kex.read(q.Super.KexId); kex == nil {
		//    --> reply 404 if not found
		result.NotFound(fmt.Errorf(`Key exchange not found`))
		goto dispatch
	}
	// -> check kex.SameSource
	if !kex.IsSameSourceString(q.Super.RemoteAddr) {
		//    --> reply 404 if !SameSource
		result.NotFound(fmt.Errorf(`Key exchange not found`))
		goto dispatch
	}
	// -> delete kex from s.kex (kex is now used)
	s.kex.remove(q.Super.KexId)
	// -> rdata = kex.DecodeAndDecrypt(data)
	if err = kex.DecodeAndDecrypt(&data, &plain); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> json.Unmarshal(rdata, &token)
	if err = json.Unmarshal(plain, &token); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> check token.UserName != `root`
	if token.UserName == `root` {
		//    --> reply 401
		result.Unauthorized(fmt.Errorf(`Cannot activate root`))
		goto dispatch
	}

	// check we have the user
	if err = s.stmt_FindUser.QueryRow(token.UserName).Scan(&userId); err == sql.ErrNoRows {
		result.Unauthorized(fmt.Errorf("Unknown user: %s", token.UserName))
		goto dispatch
	} else if err != nil {
		result.ServerError(err)
	}
	userUUID, _ = uuid.FromString(userId)

	// no account ownership verification in open mode
	if !SomaCfg.OpenInstance {
		switch s.activation {
		case `ldap`:
			if ok, err = validateLdapCredentials(token.UserName, token.Token); err != nil {
				result.ServerError(err)
				goto dispatch
			} else if !ok {
				result.Unauthorized(nil)
				goto dispatch
			}
			// fail activation if local password is the same as the
			// upstream password
			if token.Token == token.Password {
				result.Unauthorized(fmt.Errorf("User %s denied: matching local/upstream passwords", token.UserName))
				goto dispatch
			}
		case `token`: // TODO
			result.ServerError(fmt.Errorf(`Not implemented`))
			goto dispatch
		default:
			result.ServerError(fmt.Errorf("Unknown activation: %s",
				SomaCfg.Auth.Activation))
			goto dispatch
		}
	}
	// OK: validation success

	// -> scrypth64.Digest(Password, nil)
	if mcf, err = scrypth64.Digest(token.Password, nil); err != nil {
		result.Unauthorized(nil)
		goto dispatch
	}
	// -> generate token
	token.SetIPAddressString(q.Super.RemoteAddr)
	if err = token.Generate(mcf, s.key, s.seed); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	validFrom, _ = time.Parse(rfc3339Milli, token.ValidFrom)
	expiresAt, _ = time.Parse(rfc3339Milli, token.ExpiresAt)
	credExpiresAt = validFrom.Add(time.Duration(s.credExpiry) * time.Hour * 24).UTC()

	// -> open transaction
	if tx, err = s.conn.Begin(); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	defer tx.Rollback()
	// -> DB Insert: password data
	if _, err = tx.Exec(
		stmt.SetUserCredential,
		userUUID,
		mcf.String(),
		validFrom.UTC(),
		credExpiresAt.UTC(),
	); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> DB Update: activate user
	if _, err = tx.Exec(
		stmt.ActivateUser,
		userUUID,
	); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> DB Insert: token data
	if _, err = tx.Exec(
		stmt.InsertToken,
		token.Token,
		token.Salt,
		validFrom.UTC(),
		expiresAt.UTC(),
	); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> s.credentials Update
	s.credentials.insert(token.UserName, userUUID, validFrom.UTC(),
		credExpiresAt.UTC(), mcf)
	// -> s.tokens Update
	if err = s.tokens.insert(token.Token, token.ValidFrom, token.ExpiresAt,
		token.Salt); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// commit transaction
	if err = tx.Commit(); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> sdata = kex.EncryptAndEncode(&token)
	plain = []byte{}
	data = []byte{}
	if plain, err = json.Marshal(token); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	if err = kex.EncryptAndEncode(&plain, &data); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> send sdata reply
	result.Super = &msg.Supervisor{
		Verdict: 200,
		Data:    data,
	}
	result.OK()

dispatch:
	<-timer.C

conflict:
	q.Reply <- result
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Log incoming account activation requests<commit_after>/*-
Copyright (c) 2016, Jörg Pernfuß <joerg.pernfuss@1und1.de>
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/mjolnir42/scrypth64"
	"github.com/satori/go.uuid"
)

func (s *supervisor) activate_user(q *msg.Request) {
	result := msg.Result{Type: `supervisor`, Action: `activate_user`, Super: &msg.Supervisor{Action: ``}}

	var (
		timer                               *time.Timer
		plain                               []byte
		err                                 error
		kex                                 *auth.Kex
		validFrom, expiresAt, credExpiresAt time.Time
		token                               auth.Token
		userId                              string
		userUUID                            uuid.UUID
		ok                                  bool
		mcf                                 scrypth64.Mcf
		tx                                  *sql.Tx
	)
	data := q.Super.Data

	if s.readonly {
		result.Conflict(fmt.Errorf(`Readonly instance`))
		goto conflict
	}

	// start response timer
	timer = time.NewTimer(1 * time.Second)
	defer timer.Stop()

	// -> get kex
	if kex = s.kex.read(q.Super.KexId); kex == nil {
		//    --> reply 404 if not found
		result.NotFound(fmt.Errorf(`Key exchange not found`))
		goto dispatch
	}
	// -> check kex.SameSource
	if !kex.IsSameSourceString(q.Super.RemoteAddr) {
		//    --> reply 404 if !SameSource
		result.NotFound(fmt.Errorf(`Key exchange not found`))
		goto dispatch
	}
	// -> delete kex from s.kex (kex is now used)
	s.kex.remove(q.Super.KexId)
	// -> rdata = kex.DecodeAndDecrypt(data)
	if err = kex.DecodeAndDecrypt(&data, &plain); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> json.Unmarshal(rdata, &token)
	if err = json.Unmarshal(plain, &token); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// request has been decrypted, log it
	log.Printf(LogStrReq, q.Type, q.Action, token.UserName, q.Super.RemoteAddr)

	// -> check token.UserName != `root`
	if token.UserName == `root` {
		//    --> reply 401
		result.Unauthorized(fmt.Errorf(`Cannot activate root`))
		goto dispatch
	}

	// check we have the user
	if err = s.stmt_FindUser.QueryRow(token.UserName).Scan(&userId); err == sql.ErrNoRows {
		result.Unauthorized(fmt.Errorf("Unknown user: %s", token.UserName))
		goto dispatch
	} else if err != nil {
		result.ServerError(err)
	}
	userUUID, _ = uuid.FromString(userId)

	// no account ownership verification in open mode
	if !SomaCfg.OpenInstance {
		switch s.activation {
		case `ldap`:
			if ok, err = validateLdapCredentials(token.UserName, token.Token); err != nil {
				result.ServerError(err)
				goto dispatch
			} else if !ok {
				result.Unauthorized(nil)
				goto dispatch
			}
			// fail activation if local password is the same as the
			// upstream password
			if token.Token == token.Password {
				result.Unauthorized(fmt.Errorf("User %s denied: matching local/upstream passwords", token.UserName))
				goto dispatch
			}
		case `token`: // TODO
			result.ServerError(fmt.Errorf(`Not implemented`))
			goto dispatch
		default:
			result.ServerError(fmt.Errorf("Unknown activation: %s",
				SomaCfg.Auth.Activation))
			goto dispatch
		}
	}
	// OK: validation success

	// -> scrypth64.Digest(Password, nil)
	if mcf, err = scrypth64.Digest(token.Password, nil); err != nil {
		result.Unauthorized(nil)
		goto dispatch
	}
	// -> generate token
	token.SetIPAddressString(q.Super.RemoteAddr)
	if err = token.Generate(mcf, s.key, s.seed); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	validFrom, _ = time.Parse(rfc3339Milli, token.ValidFrom)
	expiresAt, _ = time.Parse(rfc3339Milli, token.ExpiresAt)
	credExpiresAt = validFrom.Add(time.Duration(s.credExpiry) * time.Hour * 24).UTC()

	// -> open transaction
	if tx, err = s.conn.Begin(); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	defer tx.Rollback()
	// -> DB Insert: password data
	if _, err = tx.Exec(
		stmt.SetUserCredential,
		userUUID,
		mcf.String(),
		validFrom.UTC(),
		credExpiresAt.UTC(),
	); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> DB Update: activate user
	if _, err = tx.Exec(
		stmt.ActivateUser,
		userUUID,
	); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> DB Insert: token data
	if _, err = tx.Exec(
		stmt.InsertToken,
		token.Token,
		token.Salt,
		validFrom.UTC(),
		expiresAt.UTC(),
	); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> s.credentials Update
	s.credentials.insert(token.UserName, userUUID, validFrom.UTC(),
		credExpiresAt.UTC(), mcf)
	// -> s.tokens Update
	if err = s.tokens.insert(token.Token, token.ValidFrom, token.ExpiresAt,
		token.Salt); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// commit transaction
	if err = tx.Commit(); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> sdata = kex.EncryptAndEncode(&token)
	plain = []byte{}
	data = []byte{}
	if plain, err = json.Marshal(token); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	if err = kex.EncryptAndEncode(&plain, &data); err != nil {
		result.ServerError(err)
		goto dispatch
	}
	// -> send sdata reply
	result.Super.Verdict = 200
	result.Super.Data = data
	result.OK()

dispatch:
	<-timer.C

conflict:
	q.Reply <- result
}

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


<commit_before>// Copyright 2010 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 smtp implements the Simple Mail Transfer Protocol as defined in RFC 5321.
// It also implements the following extensions:
//	8BITMIME  RFC 1652
//	AUTH      RFC 2554
//	STARTTLS  RFC 3207
// Additional extensions may be handled by clients.
package smtp

import (
	"crypto/tls"
	"encoding/base64"
	"errors"
	"io"
	"net"
	"net/textproto"
	"strings"
)

// A Client represents a client connection to an SMTP server.
type Client struct {
	// Text is the textproto.Conn used by the Client. It is exported to allow for
	// clients to add extensions.
	Text *textproto.Conn
	// keep a reference to the connection so it can be used to create a TLS
	// connection later
	conn net.Conn
	// whether the Client is using TLS
	tls        bool
	serverName string
	// map of supported extensions
	ext map[string]string
	// supported auth mechanisms
	auth       []string
	localName  string // the name to use in HELO/EHLO
	didHello   bool   // whether we've said HELO/EHLO
	helloError error  // the error from the hello
}

// Dial returns a new Client connected to an SMTP server at addr.
// The addr must include a port, as in "mail.example.com:smtp".
func Dial(addr string) (*Client, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}
	host, _, _ := net.SplitHostPort(addr)
	return NewClient(conn, host)
}

// NewClient returns a new Client using an existing connection and host as a
// server name to be used when authenticating.
func NewClient(conn net.Conn, host string) (*Client, error) {
	text := textproto.NewConn(conn)
	_, _, err := text.ReadResponse(220)
	if err != nil {
		text.Close()
		return nil, err
	}
	c := &Client{Text: text, conn: conn, serverName: host, localName: "localhost"}
	return c, nil
}

// Close closes the connection.
func (c *Client) Close() error {
	return c.Text.Close()
}

// hello runs a hello exchange if needed.
func (c *Client) hello() error {
	if !c.didHello {
		c.didHello = true
		err := c.ehlo()
		if err != nil {
			c.helloError = c.helo()
		}
	}
	return c.helloError
}

// Hello sends a HELO or EHLO to the server as the given host name.
// Calling this method is only necessary if the client needs control
// over the host name used. The client will introduce itself as "localhost"
// automatically otherwise. If Hello is called, it must be called before
// any of the other methods.
func (c *Client) Hello(localName string) error {
	if c.didHello {
		return errors.New("smtp: Hello called after other methods")
	}
	c.localName = localName
	return c.hello()
}

// cmd is a convenience function that sends a command and returns the response
func (c *Client) cmd(expectCode int, format string, args ...interface{}) (int, string, error) {
	id, err := c.Text.Cmd(format, args...)
	if err != nil {
		return 0, "", err
	}
	c.Text.StartResponse(id)
	defer c.Text.EndResponse(id)
	code, msg, err := c.Text.ReadResponse(expectCode)
	return code, msg, err
}

// helo sends the HELO greeting to the server. It should be used only when the
// server does not support ehlo.
func (c *Client) helo() error {
	c.ext = nil
	_, _, err := c.cmd(250, "HELO %s", c.localName)
	return err
}

// ehlo sends the EHLO (extended hello) greeting to the server. It
// should be the preferred greeting for servers that support it.
func (c *Client) ehlo() error {
	_, msg, err := c.cmd(250, "EHLO %s", c.localName)
	if err != nil {
		return err
	}
	ext := make(map[string]string)
	extList := strings.Split(msg, "\n")
	if len(extList) > 1 {
		extList = extList[1:]
		for _, line := range extList {
			args := strings.SplitN(line, " ", 2)
			if len(args) > 1 {
				ext[args[0]] = args[1]
			} else {
				ext[args[0]] = ""
			}
		}
	}
	if mechs, ok := ext["AUTH"]; ok {
		c.auth = strings.Split(mechs, " ")
	}
	c.ext = ext
	return err
}

// StartTLS sends the STARTTLS command and encrypts all further communication.
// Only servers that advertise the STARTTLS extension support this function.
func (c *Client) StartTLS(config *tls.Config) error {
	if err := c.hello(); err != nil {
		return err
	}
	_, _, err := c.cmd(220, "STARTTLS")
	if err != nil {
		return err
	}
	c.conn = tls.Client(c.conn, config)
	c.Text = textproto.NewConn(c.conn)
	c.tls = true
	return c.ehlo()
}

// TLSConnectionState returns the client's TLS connection state.
// The return values are their zero values if StartTLS did
// not succeed.
func (c *Client) TLSConnectionState() (state tls.ConnectionState, ok bool) {
	tc, ok := c.conn.(*tls.Conn)
	if !ok {
		return
	}
	return tc.ConnectionState(), true
}

// Verify checks the validity of an email address on the server.
// If Verify returns nil, the address is valid. A non-nil return
// does not necessarily indicate an invalid address. Many servers
// will not verify addresses for security reasons.
func (c *Client) Verify(addr string) error {
	if err := c.hello(); err != nil {
		return err
	}
	_, _, err := c.cmd(250, "VRFY %s", addr)
	return err
}

// Auth authenticates a client using the provided authentication mechanism.
// A failed authentication closes the connection.
// Only servers that advertise the AUTH extension support this function.
func (c *Client) Auth(a Auth) error {
	if err := c.hello(); err != nil {
		return err
	}
	encoding := base64.StdEncoding
	mech, resp, err := a.Start(&ServerInfo{c.serverName, c.tls, c.auth})
	if err != nil {
		c.Quit()
		return err
	}
	resp64 := make([]byte, encoding.EncodedLen(len(resp)))
	encoding.Encode(resp64, resp)
	code, msg64, err := c.cmd(0, "AUTH %s %s", mech, resp64)
	for err == nil {
		var msg []byte
		switch code {
		case 334:
			msg, err = encoding.DecodeString(msg64)
		case 235:
			// the last message isn't base64 because it isn't a challenge
			msg = []byte(msg64)
		default:
			err = &textproto.Error{Code: code, Msg: msg64}
		}
		if err == nil {
			resp, err = a.Next(msg, code == 334)
		}
		if err != nil {
			// abort the AUTH
			c.cmd(501, "*")
			c.Quit()
			break
		}
		if resp == nil {
			break
		}
		resp64 = make([]byte, encoding.EncodedLen(len(resp)))
		encoding.Encode(resp64, resp)
		code, msg64, err = c.cmd(0, string(resp64))
	}
	return err
}

// Mail issues a MAIL command to the server using the provided email address.
// If the server supports the 8BITMIME extension, Mail adds the BODY=8BITMIME
// parameter.
// This initiates a mail transaction and is followed by one or more Rcpt calls.
func (c *Client) Mail(from string) error {
	if err := c.hello(); err != nil {
		return err
	}
	cmdStr := "MAIL FROM:<%s>"
	if c.ext != nil {
		if _, ok := c.ext["8BITMIME"]; ok {
			cmdStr += " BODY=8BITMIME"
		}
	}
	_, _, err := c.cmd(250, cmdStr, from)
	return err
}

// Rcpt issues a RCPT command to the server using the provided email address.
// A call to Rcpt must be preceded by a call to Mail and may be followed by
// a Data call or another Rcpt call.
func (c *Client) Rcpt(to string) error {
	_, _, err := c.cmd(25, "RCPT TO:<%s>", to)
	return err
}

type dataCloser struct {
	c *Client
	io.WriteCloser
}

func (d *dataCloser) Close() error {
	d.WriteCloser.Close()
	_, _, err := d.c.Text.ReadResponse(250)
	return err
}

// Data issues a DATA command to the server and returns a writer that
// can be used to write the mail headers and body. The caller should
// close the writer before calling any more methods on c. A call to
// Data must be preceded by one or more calls to Rcpt.
func (c *Client) Data() (io.WriteCloser, error) {
	_, _, err := c.cmd(354, "DATA")
	if err != nil {
		return nil, err
	}
	return &dataCloser{c, c.Text.DotWriter()}, nil
}

var testHookStartTLS func(*tls.Config) // nil, except for tests

// SendMail connects to the server at addr, switches to TLS if
// possible, authenticates with the optional mechanism a if possible,
// and then sends an email from address from, to addresses to, with
// message msg.
// The addr must include a port, as in "mail.example.com:smtp".
//
// The addresses in the to parameter are the SMTP RCPT addresses.
//
// The msg parameter should be an RFC 822-style email with headers
// first, a blank line, and then the message body. The lines of msg
// should be CRLF terminated. The msg headers should usually include
// fields such as "From", "To", "Subject", and "Cc".  Sending "Bcc"
// messages is accomplished by including an email address in the to
// parameter but not including it in the msg headers.
//
// The SendMail function and the the net/smtp package are low-level
// mechanisms and provide no support for DKIM signing, MIME
// attachments (see the mime/multipart package), or other mail
// functionality. Higher-level packages exist outside of the standard
// library.
func SendMail(addr string, a Auth, from string, to []string, msg []byte) error {
	c, err := Dial(addr)
	if err != nil {
		return err
	}
	defer c.Close()
	if err = c.hello(); err != nil {
		return err
	}
	if ok, _ := c.Extension("STARTTLS"); ok {
		config := &tls.Config{ServerName: c.serverName}
		if testHookStartTLS != nil {
			testHookStartTLS(config)
		}
		if err = c.StartTLS(config); err != nil {
			return err
		}
	}
	if a != nil && c.ext != nil {
		if _, ok := c.ext["AUTH"]; ok {
			if err = c.Auth(a); err != nil {
				return err
			}
		}
	}
	if err = c.Mail(from); err != nil {
		return err
	}
	for _, addr := range to {
		if err = c.Rcpt(addr); err != nil {
			return err
		}
	}
	w, err := c.Data()
	if err != nil {
		return err
	}
	_, err = w.Write(msg)
	if err != nil {
		return err
	}
	err = w.Close()
	if err != nil {
		return err
	}
	return c.Quit()
}

// Extension reports whether an extension is support by the server.
// The extension name is case-insensitive. If the extension is supported,
// Extension also returns a string that contains any parameters the
// server specifies for the extension.
func (c *Client) Extension(ext string) (bool, string) {
	if err := c.hello(); err != nil {
		return false, ""
	}
	if c.ext == nil {
		return false, ""
	}
	ext = strings.ToUpper(ext)
	param, ok := c.ext[ext]
	return ok, param
}

// Reset sends the RSET command to the server, aborting the current mail
// transaction.
func (c *Client) Reset() error {
	if err := c.hello(); err != nil {
		return err
	}
	_, _, err := c.cmd(250, "RSET")
	return err
}

// Quit sends the QUIT command and closes the connection to the server.
func (c *Client) Quit() error {
	if err := c.hello(); err != nil {
		return err
	}
	_, _, err := c.cmd(221, "QUIT")
	if err != nil {
		return err
	}
	return c.Text.Close()
}
<commit_msg>net/smtp: document that the smtp package is frozen<commit_after>// Copyright 2010 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 smtp implements the Simple Mail Transfer Protocol as defined in RFC 5321.
// It also implements the following extensions:
//	8BITMIME  RFC 1652
//	AUTH      RFC 2554
//	STARTTLS  RFC 3207
// Additional extensions may be handled by clients.
//
// The smtp package is frozen and not accepting new features.
// Some external packages provide more functionality. See:
//
//   https://godoc.org/?q=smtp
package smtp

import (
	"crypto/tls"
	"encoding/base64"
	"errors"
	"io"
	"net"
	"net/textproto"
	"strings"
)

// A Client represents a client connection to an SMTP server.
type Client struct {
	// Text is the textproto.Conn used by the Client. It is exported to allow for
	// clients to add extensions.
	Text *textproto.Conn
	// keep a reference to the connection so it can be used to create a TLS
	// connection later
	conn net.Conn
	// whether the Client is using TLS
	tls        bool
	serverName string
	// map of supported extensions
	ext map[string]string
	// supported auth mechanisms
	auth       []string
	localName  string // the name to use in HELO/EHLO
	didHello   bool   // whether we've said HELO/EHLO
	helloError error  // the error from the hello
}

// Dial returns a new Client connected to an SMTP server at addr.
// The addr must include a port, as in "mail.example.com:smtp".
func Dial(addr string) (*Client, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}
	host, _, _ := net.SplitHostPort(addr)
	return NewClient(conn, host)
}

// NewClient returns a new Client using an existing connection and host as a
// server name to be used when authenticating.
func NewClient(conn net.Conn, host string) (*Client, error) {
	text := textproto.NewConn(conn)
	_, _, err := text.ReadResponse(220)
	if err != nil {
		text.Close()
		return nil, err
	}
	c := &Client{Text: text, conn: conn, serverName: host, localName: "localhost"}
	return c, nil
}

// Close closes the connection.
func (c *Client) Close() error {
	return c.Text.Close()
}

// hello runs a hello exchange if needed.
func (c *Client) hello() error {
	if !c.didHello {
		c.didHello = true
		err := c.ehlo()
		if err != nil {
			c.helloError = c.helo()
		}
	}
	return c.helloError
}

// Hello sends a HELO or EHLO to the server as the given host name.
// Calling this method is only necessary if the client needs control
// over the host name used. The client will introduce itself as "localhost"
// automatically otherwise. If Hello is called, it must be called before
// any of the other methods.
func (c *Client) Hello(localName string) error {
	if c.didHello {
		return errors.New("smtp: Hello called after other methods")
	}
	c.localName = localName
	return c.hello()
}

// cmd is a convenience function that sends a command and returns the response
func (c *Client) cmd(expectCode int, format string, args ...interface{}) (int, string, error) {
	id, err := c.Text.Cmd(format, args...)
	if err != nil {
		return 0, "", err
	}
	c.Text.StartResponse(id)
	defer c.Text.EndResponse(id)
	code, msg, err := c.Text.ReadResponse(expectCode)
	return code, msg, err
}

// helo sends the HELO greeting to the server. It should be used only when the
// server does not support ehlo.
func (c *Client) helo() error {
	c.ext = nil
	_, _, err := c.cmd(250, "HELO %s", c.localName)
	return err
}

// ehlo sends the EHLO (extended hello) greeting to the server. It
// should be the preferred greeting for servers that support it.
func (c *Client) ehlo() error {
	_, msg, err := c.cmd(250, "EHLO %s", c.localName)
	if err != nil {
		return err
	}
	ext := make(map[string]string)
	extList := strings.Split(msg, "\n")
	if len(extList) > 1 {
		extList = extList[1:]
		for _, line := range extList {
			args := strings.SplitN(line, " ", 2)
			if len(args) > 1 {
				ext[args[0]] = args[1]
			} else {
				ext[args[0]] = ""
			}
		}
	}
	if mechs, ok := ext["AUTH"]; ok {
		c.auth = strings.Split(mechs, " ")
	}
	c.ext = ext
	return err
}

// StartTLS sends the STARTTLS command and encrypts all further communication.
// Only servers that advertise the STARTTLS extension support this function.
func (c *Client) StartTLS(config *tls.Config) error {
	if err := c.hello(); err != nil {
		return err
	}
	_, _, err := c.cmd(220, "STARTTLS")
	if err != nil {
		return err
	}
	c.conn = tls.Client(c.conn, config)
	c.Text = textproto.NewConn(c.conn)
	c.tls = true
	return c.ehlo()
}

// TLSConnectionState returns the client's TLS connection state.
// The return values are their zero values if StartTLS did
// not succeed.
func (c *Client) TLSConnectionState() (state tls.ConnectionState, ok bool) {
	tc, ok := c.conn.(*tls.Conn)
	if !ok {
		return
	}
	return tc.ConnectionState(), true
}

// Verify checks the validity of an email address on the server.
// If Verify returns nil, the address is valid. A non-nil return
// does not necessarily indicate an invalid address. Many servers
// will not verify addresses for security reasons.
func (c *Client) Verify(addr string) error {
	if err := c.hello(); err != nil {
		return err
	}
	_, _, err := c.cmd(250, "VRFY %s", addr)
	return err
}

// Auth authenticates a client using the provided authentication mechanism.
// A failed authentication closes the connection.
// Only servers that advertise the AUTH extension support this function.
func (c *Client) Auth(a Auth) error {
	if err := c.hello(); err != nil {
		return err
	}
	encoding := base64.StdEncoding
	mech, resp, err := a.Start(&ServerInfo{c.serverName, c.tls, c.auth})
	if err != nil {
		c.Quit()
		return err
	}
	resp64 := make([]byte, encoding.EncodedLen(len(resp)))
	encoding.Encode(resp64, resp)
	code, msg64, err := c.cmd(0, "AUTH %s %s", mech, resp64)
	for err == nil {
		var msg []byte
		switch code {
		case 334:
			msg, err = encoding.DecodeString(msg64)
		case 235:
			// the last message isn't base64 because it isn't a challenge
			msg = []byte(msg64)
		default:
			err = &textproto.Error{Code: code, Msg: msg64}
		}
		if err == nil {
			resp, err = a.Next(msg, code == 334)
		}
		if err != nil {
			// abort the AUTH
			c.cmd(501, "*")
			c.Quit()
			break
		}
		if resp == nil {
			break
		}
		resp64 = make([]byte, encoding.EncodedLen(len(resp)))
		encoding.Encode(resp64, resp)
		code, msg64, err = c.cmd(0, string(resp64))
	}
	return err
}

// Mail issues a MAIL command to the server using the provided email address.
// If the server supports the 8BITMIME extension, Mail adds the BODY=8BITMIME
// parameter.
// This initiates a mail transaction and is followed by one or more Rcpt calls.
func (c *Client) Mail(from string) error {
	if err := c.hello(); err != nil {
		return err
	}
	cmdStr := "MAIL FROM:<%s>"
	if c.ext != nil {
		if _, ok := c.ext["8BITMIME"]; ok {
			cmdStr += " BODY=8BITMIME"
		}
	}
	_, _, err := c.cmd(250, cmdStr, from)
	return err
}

// Rcpt issues a RCPT command to the server using the provided email address.
// A call to Rcpt must be preceded by a call to Mail and may be followed by
// a Data call or another Rcpt call.
func (c *Client) Rcpt(to string) error {
	_, _, err := c.cmd(25, "RCPT TO:<%s>", to)
	return err
}

type dataCloser struct {
	c *Client
	io.WriteCloser
}

func (d *dataCloser) Close() error {
	d.WriteCloser.Close()
	_, _, err := d.c.Text.ReadResponse(250)
	return err
}

// Data issues a DATA command to the server and returns a writer that
// can be used to write the mail headers and body. The caller should
// close the writer before calling any more methods on c. A call to
// Data must be preceded by one or more calls to Rcpt.
func (c *Client) Data() (io.WriteCloser, error) {
	_, _, err := c.cmd(354, "DATA")
	if err != nil {
		return nil, err
	}
	return &dataCloser{c, c.Text.DotWriter()}, nil
}

var testHookStartTLS func(*tls.Config) // nil, except for tests

// SendMail connects to the server at addr, switches to TLS if
// possible, authenticates with the optional mechanism a if possible,
// and then sends an email from address from, to addresses to, with
// message msg.
// The addr must include a port, as in "mail.example.com:smtp".
//
// The addresses in the to parameter are the SMTP RCPT addresses.
//
// The msg parameter should be an RFC 822-style email with headers
// first, a blank line, and then the message body. The lines of msg
// should be CRLF terminated. The msg headers should usually include
// fields such as "From", "To", "Subject", and "Cc".  Sending "Bcc"
// messages is accomplished by including an email address in the to
// parameter but not including it in the msg headers.
//
// The SendMail function and the the net/smtp package are low-level
// mechanisms and provide no support for DKIM signing, MIME
// attachments (see the mime/multipart package), or other mail
// functionality. Higher-level packages exist outside of the standard
// library.
func SendMail(addr string, a Auth, from string, to []string, msg []byte) error {
	c, err := Dial(addr)
	if err != nil {
		return err
	}
	defer c.Close()
	if err = c.hello(); err != nil {
		return err
	}
	if ok, _ := c.Extension("STARTTLS"); ok {
		config := &tls.Config{ServerName: c.serverName}
		if testHookStartTLS != nil {
			testHookStartTLS(config)
		}
		if err = c.StartTLS(config); err != nil {
			return err
		}
	}
	if a != nil && c.ext != nil {
		if _, ok := c.ext["AUTH"]; ok {
			if err = c.Auth(a); err != nil {
				return err
			}
		}
	}
	if err = c.Mail(from); err != nil {
		return err
	}
	for _, addr := range to {
		if err = c.Rcpt(addr); err != nil {
			return err
		}
	}
	w, err := c.Data()
	if err != nil {
		return err
	}
	_, err = w.Write(msg)
	if err != nil {
		return err
	}
	err = w.Close()
	if err != nil {
		return err
	}
	return c.Quit()
}

// Extension reports whether an extension is support by the server.
// The extension name is case-insensitive. If the extension is supported,
// Extension also returns a string that contains any parameters the
// server specifies for the extension.
func (c *Client) Extension(ext string) (bool, string) {
	if err := c.hello(); err != nil {
		return false, ""
	}
	if c.ext == nil {
		return false, ""
	}
	ext = strings.ToUpper(ext)
	param, ok := c.ext[ext]
	return ok, param
}

// Reset sends the RSET command to the server, aborting the current mail
// transaction.
func (c *Client) Reset() error {
	if err := c.hello(); err != nil {
		return err
	}
	_, _, err := c.cmd(250, "RSET")
	return err
}

// Quit sends the QUIT command and closes the connection to the server.
func (c *Client) Quit() error {
	if err := c.hello(); err != nil {
		return err
	}
	_, _, err := c.cmd(221, "QUIT")
	if err != nil {
		return err
	}
	return c.Text.Close()
}
<|endoftext|>


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

package gb

import (
	"errors"
	"fmt"
	"go/ast"
	"go/build"
	"go/doc"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strings"
	"sync"
)

func FindGbProject(path string) (string, error) {
	if path == "" {
		return "", fmt.Errorf("project root is blank")
	}
	start := path
	for path != filepath.Dir(path) {
		root := filepath.Join(path, "src")
		if _, err := os.Stat(root); err != nil {
			if os.IsNotExist(err) {
				path = filepath.Dir(path)
				continue
			}
			return "", err
		}
		return path, nil
	}

	return "", fmt.Errorf(`could not find project root in "%s" or its parents`, start)
}

// GoPath returns a GOPATH for path p.
func GoPath(p string) string {
	goPath := os.Getenv("GOPATH")
	r := runtime.GOROOT()
	if r != "" {
		goPath = goPath + string(filepath.ListSeparator) + r
	}

	p = filepath.Clean(p)

	for _, root := range filepath.SplitList(goPath) {
		if strings.HasPrefix(p, filepath.Join(root, "src")+string(filepath.Separator)) {
			return goPath
		}
	}

	project, err := FindGbProject(p)
	if err == nil {
		parent, child := filepath.Split(project)
		if child == "vendor" {
			project = parent[:len(parent)-1]
		}
		return project + string(filepath.ListSeparator) +
			filepath.Join(project, "vendor") + string(filepath.ListSeparator) + goPath
	}

	return goPath
}

var goBuildDefaultMu sync.Mutex

// WithGoBuildForPath sets the go/build Default.GOPATH to GoPath(p) under a
// mutex. The returned function restores Default.GOPATH to its original value
// and unlocks the mutex.
//
// This function intended to be used to the golang.org/x/tools/imports and
// other packages that use go/build Default.
func WithGoBuildForPath(p string) func() {
	dir, _ := filepath.Split(p)
	goBuildDefaultMu.Lock()
	original := build.Default.GOPATH
	build.Default.GOPATH = GoPath(dir)
	return func() {
		build.Default.GOPATH = original
		goBuildDefaultMu.Unlock()
	}
}

// Package represents a Go package.
type Package struct {
	FSet     *token.FileSet
	Build    *build.Package
	AST      *ast.Package
	Doc      *doc.Package
	Examples []*doc.Example
	Errors   []error
}

// Flags for LoadPackage.
const (
	LoadDoc = 1 << iota
	LoadExamples
	LoadUnexported
)

// LoadPackage Import returns details about the Go package named by the import
// path, interpreting local import paths relative to the srcDir directory.
func LoadPackage(importPath string, srcDir string, flags int) (*Package, error) {
	bpkg, err := build.Default.Import(importPath, srcDir, build.ImportComment)
	if _, ok := err.(*build.NoGoError); ok {
		return &Package{Build: bpkg}, nil
	}
	if err != nil {
		return nil, err
	}

	pkg := &Package{
		FSet:  token.NewFileSet(),
		Build: bpkg,
	}

	files := make(map[string]*ast.File)
	for _, name := range append(pkg.Build.GoFiles, pkg.Build.CgoFiles...) {
		file, err := pkg.parseFile(name)
		if err != nil {
			pkg.Errors = append(pkg.Errors, err)
			continue
		}
		files[name] = file
	}

	pkg.AST, _ = ast.NewPackage(pkg.FSet, files, simpleImporter, nil)

	if flags&LoadDoc != 0 {
		mode := doc.Mode(0)
		if pkg.Build.ImportPath == "builtin" || flags&LoadUnexported != 0 {
			mode |= doc.AllDecls
		}
		pkg.Doc = doc.New(pkg.AST, pkg.Build.ImportPath, mode)
		if pkg.Build.ImportPath == "builtin" {
			for _, t := range pkg.Doc.Types {
				pkg.Doc.Funcs = append(pkg.Doc.Funcs, t.Funcs...)
				t.Funcs = nil
			}
			sort.Sort(byFuncName(pkg.Doc.Funcs))
		}
	}

	if flags&LoadExamples != 0 {
		for _, name := range append(pkg.Build.TestGoFiles, pkg.Build.XTestGoFiles...) {
			file, err := pkg.parseFile(name)
			if err != nil {
				pkg.Errors = append(pkg.Errors, err)
				continue
			}
			pkg.Examples = append(pkg.Examples, doc.Examples(file)...)
		}
	}

	return pkg, nil
}

type byFuncName []*doc.Func

func (s byFuncName) Len() int           { return len(s) }
func (s byFuncName) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s byFuncName) Less(i, j int) bool { return s[i].Name < s[j].Name }

func (pkg *Package) parseFile(name string) (*ast.File, error) {
	p, err := ioutil.ReadFile(filepath.Join(pkg.Build.Dir, name))
	if err != nil {
		return nil, err
	}
	// overwrite //line comments
	for _, m := range linePat.FindAllIndex(p, -1) {
		for i := m[0] + 2; i < m[1]; i++ {
			p[i] = ' '
		}
	}
	return parser.ParseFile(pkg.FSet, name, p, parser.ParseComments)
}

var linePat = regexp.MustCompile(`(?m)^//line .*$`)

func simpleImporter(imports map[string]*ast.Object, path string) (*ast.Object, error) {
	pkg := imports[path]
	if pkg != nil {
		return pkg, nil
	}

	n := GuessPackageNameFromPath(path)
	if n == "" {
		return nil, errors.New("package not found")
	}

	pkg = ast.NewObj(ast.Pkg, n)
	pkg.Data = ast.NewScope(nil)
	imports[path] = pkg
	return pkg, nil
}

var packageNamePats = []*regexp.Regexp{
	// Last element with .suffix removed.
	regexp.MustCompile(`/([^-./]+)[-.](?:git|svn|hg|bzr|v\d+)$`),

	// Last element with "go" prefix or suffix removed.
	regexp.MustCompile(`/([^-./]+)[-.]go$`),
	regexp.MustCompile(`/go[-.]([^-./]+)$`),

	// It's also common for the last element of the path to contain an
	// extra "go" prefix, but not always. TODO: examine unresolved ids to
	// detect when trimming the "go" prefix is appropriate.

	// Last component of path.
	regexp.MustCompile(`([^/]+)$`),
}

// GuessPackageNameFromPath guesses the package name from the package path.
func GuessPackageNameFromPath(path string) string {
	// Guess the package name without importing it.
	for _, pat := range packageNamePats {
		m := pat.FindStringSubmatch(path)
		if m != nil {
			return m[1]
		}
	}
	return ""
}
<commit_msg>Fix WithGoBuildForPath to based by dir & Add UseAllFiles to false<commit_after>// Copyright 2015 Gary Burd. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gb

import (
	"errors"
	"fmt"
	"go/ast"
	"go/build"
	"go/doc"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strings"
	"sync"
)

func FindGbProject(path string) (string, error) {
	if path == "" {
		return "", fmt.Errorf("project root is blank")
	}
	start := path
	for path != filepath.Dir(path) {
		root := filepath.Join(path, "src")
		if _, err := os.Stat(root); err != nil {
			if os.IsNotExist(err) {
				path = filepath.Dir(path)
				continue
			}
			return "", err
		}
		return path, nil
	}

	return "", fmt.Errorf(`could not find project root in "%s" or its parents`, start)
}

// GoPath returns a GOPATH for path p.
func GoPath(p string) string {
	goPath := os.Getenv("GOPATH")
	r := runtime.GOROOT()
	if r != "" {
		goPath = goPath + string(filepath.ListSeparator) + r
	}

	p = filepath.Clean(p)

	for _, root := range filepath.SplitList(goPath) {
		if strings.HasPrefix(p, filepath.Join(root, "src")+string(filepath.Separator)) {
			return goPath
		}
	}

	project, err := FindGbProject(p)
	if err == nil {
		parent, child := filepath.Split(project)
		if child == "vendor" {
			project = parent[:len(parent)-1]
		}
		return project + string(filepath.ListSeparator) +
			filepath.Join(project, "vendor") + string(filepath.ListSeparator) + goPath
	}

	return goPath
}

var goBuildDefaultMu sync.Mutex

// WithGoBuildForPath sets the go/build Default.GOPATH to GoPath(p) under a
// mutex. The returned function restores Default.GOPATH to its original value
// and unlocks the mutex.
//
// This function intended to be used to the golang.org/x/tools/imports and
// other packages that use go/build Default.
func WithGoBuildForPath(p string) func() {
	goBuildDefaultMu.Lock()
	original := build.Default.GOPATH
	build.Default.GOPATH = GoPath(p)
	build.Default.UseAllFiles = false
	return func() {
		build.Default.GOPATH = original
		goBuildDefaultMu.Unlock()
	}
}

// Package represents a Go package.
type Package struct {
	FSet     *token.FileSet
	Build    *build.Package
	AST      *ast.Package
	Doc      *doc.Package
	Examples []*doc.Example
	Errors   []error
}

// Flags for LoadPackage.
const (
	LoadDoc = 1 << iota
	LoadExamples
	LoadUnexported
)

// LoadPackage Import returns details about the Go package named by the import
// path, interpreting local import paths relative to the srcDir directory.
func LoadPackage(importPath string, srcDir string, flags int) (*Package, error) {
	bpkg, err := build.Default.Import(importPath, srcDir, build.ImportComment)
	if _, ok := err.(*build.NoGoError); ok {
		return &Package{Build: bpkg}, nil
	}
	if err != nil {
		return nil, err
	}

	pkg := &Package{
		FSet:  token.NewFileSet(),
		Build: bpkg,
	}

	files := make(map[string]*ast.File)
	for _, name := range append(pkg.Build.GoFiles, pkg.Build.CgoFiles...) {
		file, err := pkg.parseFile(name)
		if err != nil {
			pkg.Errors = append(pkg.Errors, err)
			continue
		}
		files[name] = file
	}

	pkg.AST, _ = ast.NewPackage(pkg.FSet, files, simpleImporter, nil)

	if flags&LoadDoc != 0 {
		mode := doc.Mode(0)
		if pkg.Build.ImportPath == "builtin" || flags&LoadUnexported != 0 {
			mode |= doc.AllDecls
		}
		pkg.Doc = doc.New(pkg.AST, pkg.Build.ImportPath, mode)
		if pkg.Build.ImportPath == "builtin" {
			for _, t := range pkg.Doc.Types {
				pkg.Doc.Funcs = append(pkg.Doc.Funcs, t.Funcs...)
				t.Funcs = nil
			}
			sort.Sort(byFuncName(pkg.Doc.Funcs))
		}
	}

	if flags&LoadExamples != 0 {
		for _, name := range append(pkg.Build.TestGoFiles, pkg.Build.XTestGoFiles...) {
			file, err := pkg.parseFile(name)
			if err != nil {
				pkg.Errors = append(pkg.Errors, err)
				continue
			}
			pkg.Examples = append(pkg.Examples, doc.Examples(file)...)
		}
	}

	return pkg, nil
}

type byFuncName []*doc.Func

func (s byFuncName) Len() int           { return len(s) }
func (s byFuncName) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s byFuncName) Less(i, j int) bool { return s[i].Name < s[j].Name }

func (pkg *Package) parseFile(name string) (*ast.File, error) {
	p, err := ioutil.ReadFile(filepath.Join(pkg.Build.Dir, name))
	if err != nil {
		return nil, err
	}
	// overwrite //line comments
	for _, m := range linePat.FindAllIndex(p, -1) {
		for i := m[0] + 2; i < m[1]; i++ {
			p[i] = ' '
		}
	}
	return parser.ParseFile(pkg.FSet, name, p, parser.ParseComments)
}

var linePat = regexp.MustCompile(`(?m)^//line .*$`)

func simpleImporter(imports map[string]*ast.Object, path string) (*ast.Object, error) {
	pkg := imports[path]
	if pkg != nil {
		return pkg, nil
	}

	n := GuessPackageNameFromPath(path)
	if n == "" {
		return nil, errors.New("package not found")
	}

	pkg = ast.NewObj(ast.Pkg, n)
	pkg.Data = ast.NewScope(nil)
	imports[path] = pkg
	return pkg, nil
}

var packageNamePats = []*regexp.Regexp{
	// Last element with .suffix removed.
	regexp.MustCompile(`/([^-./]+)[-.](?:git|svn|hg|bzr|v\d+)$`),

	// Last element with "go" prefix or suffix removed.
	regexp.MustCompile(`/([^-./]+)[-.]go$`),
	regexp.MustCompile(`/go[-.]([^-./]+)$`),

	// It's also common for the last element of the path to contain an
	// extra "go" prefix, but not always. TODO: examine unresolved ids to
	// detect when trimming the "go" prefix is appropriate.

	// Last component of path.
	regexp.MustCompile(`([^/]+)$`),
}

// GuessPackageNameFromPath guesses the package name from the package path.
func GuessPackageNameFromPath(path string) string {
	// Guess the package name without importing it.
	for _, pat := range packageNamePats {
		m := pat.FindStringSubmatch(path)
		if m != nil {
			return m[1]
		}
	}
	return ""
}
<|endoftext|>