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

package main

import (
	"bufio"
	"cmd/internal/browser"
	"flag"
	"fmt"
	"html/template"
	"internal/trace"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"runtime"
	"runtime/debug"
	"sync"

	_ "net/http/pprof" // Required to use pprof
)

const usageMessage = "" +
	`Usage of 'go tool trace':
Given a trace file produced by 'go test':
	go test -trace=trace.out pkg

Open a web browser displaying trace:
	go tool trace [flags] [pkg.test] trace.out

Generate a pprof-like profile from the trace:
    go tool trace -pprof=TYPE [pkg.test] trace.out

[pkg.test] argument is required for traces produced by Go 1.6 and below.
Go 1.7 does not require the binary argument.

Supported profile types are:
    - net: network blocking profile
    - sync: synchronization blocking profile
    - syscall: syscall blocking profile
    - sched: scheduler latency profile

Flags:
	-http=addr: HTTP service address (e.g., ':6060')
	-pprof=type: print a pprof-like profile instead
	-d: print debug info such as parsed events

Note that while the various profiles available when launching
'go tool trace' work on every browser, the trace viewer itself
(the 'view trace' page) comes from the Chrome/Chromium project
and is only actively tested on that browser.
`

var (
	httpFlag  = flag.String("http", "localhost:0", "HTTP service address (e.g., ':6060')")
	pprofFlag = flag.String("pprof", "", "print a pprof-like profile instead")
	debugFlag = flag.Bool("d", false, "print debug information such as parsed events list")

	// The binary file name, left here for serveSVGProfile.
	programBinary string
	traceFile     string
)

func main() {
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, usageMessage)
		os.Exit(2)
	}
	flag.Parse()

	// Go 1.7 traces embed symbol info and does not require the binary.
	// But we optionally accept binary as first arg for Go 1.5 traces.
	switch flag.NArg() {
	case 1:
		traceFile = flag.Arg(0)
	case 2:
		programBinary = flag.Arg(0)
		traceFile = flag.Arg(1)
	default:
		flag.Usage()
	}

	var pprofFunc func(io.Writer, *http.Request) error
	switch *pprofFlag {
	case "net":
		pprofFunc = pprofByGoroutine(computePprofIO)
	case "sync":
		pprofFunc = pprofByGoroutine(computePprofBlock)
	case "syscall":
		pprofFunc = pprofByGoroutine(computePprofSyscall)
	case "sched":
		pprofFunc = pprofByGoroutine(computePprofSched)
	}
	if pprofFunc != nil {
		if err := pprofFunc(os.Stdout, &http.Request{}); err != nil {
			dief("failed to generate pprof: %v\n", err)
		}
		os.Exit(0)
	}
	if *pprofFlag != "" {
		dief("unknown pprof type %s\n", *pprofFlag)
	}

	ln, err := net.Listen("tcp", *httpFlag)
	if err != nil {
		dief("failed to create server socket: %v\n", err)
	}

	log.Print("Parsing trace...")
	res, err := parseTrace()
	if err != nil {
		dief("%v\n", err)
	}

	if *debugFlag {
		trace.Print(res.Events)
		os.Exit(0)
	}
	reportMemoryUsage("after parsing trace")
	debug.FreeOSMemory()

	log.Print("Splitting trace...")
	ranges = splitTrace(res)
	reportMemoryUsage("after spliting trace")
	debug.FreeOSMemory()

	addr := "http://" + ln.Addr().String()
	log.Printf("Opening browser. Trace viewer is listening on %s", addr)
	browser.Open(addr)

	// Start http server.
	http.HandleFunc("/", httpMain)
	err = http.Serve(ln, nil)
	dief("failed to start http server: %v\n", err)
}

var ranges []Range

var loader struct {
	once sync.Once
	res  trace.ParseResult
	err  error
}

// parseEvents is a compatibility wrapper that returns only
// the Events part of trace.ParseResult returned by parseTrace.
func parseEvents() ([]*trace.Event, error) {
	res, err := parseTrace()
	if err != nil {
		return nil, err
	}
	return res.Events, err
}

func parseTrace() (trace.ParseResult, error) {
	loader.once.Do(func() {
		tracef, err := os.Open(traceFile)
		if err != nil {
			loader.err = fmt.Errorf("failed to open trace file: %v", err)
			return
		}
		defer tracef.Close()

		// Parse and symbolize.
		res, err := trace.Parse(bufio.NewReader(tracef), programBinary)
		if err != nil {
			loader.err = fmt.Errorf("failed to parse trace: %v", err)
			return
		}
		loader.res = res
	})
	return loader.res, loader.err
}

// httpMain serves the starting page.
func httpMain(w http.ResponseWriter, r *http.Request) {
	if err := templMain.Execute(w, ranges); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

var templMain = template.Must(template.New("").Parse(`

{{if $}}
	{{range $e := $}}
		View trace ({{$e.Name}})

	{{end}}

{{else}}
	View trace

{{end}}
Goroutine analysis

Network blocking profile (⬇)

Synchronization blocking profile (⬇)

Syscall blocking profile (⬇)

Scheduler latency profile (⬇)

User-defined tasks

User-defined regions

Minimum mutator utilization

`))

func dief(msg string, args ...any) {
	fmt.Fprintf(os.Stderr, msg, args...)
	os.Exit(1)
}

var debugMemoryUsage bool

func init() {
	v := os.Getenv("DEBUG_MEMORY_USAGE")
	debugMemoryUsage = v != ""
}

func reportMemoryUsage(msg string) {
	if !debugMemoryUsage {
		return
	}
	var s runtime.MemStats
	runtime.ReadMemStats(&s)
	w := os.Stderr
	fmt.Fprintf(w, "%s\n", msg)
	fmt.Fprintf(w, " Alloc:\t%d Bytes\n", s.Alloc)
	fmt.Fprintf(w, " Sys:\t%d Bytes\n", s.Sys)
	fmt.Fprintf(w, " HeapReleased:\t%d Bytes\n", s.HeapReleased)
	fmt.Fprintf(w, " HeapSys:\t%d Bytes\n", s.HeapSys)
	fmt.Fprintf(w, " HeapInUse:\t%d Bytes\n", s.HeapInuse)
	fmt.Fprintf(w, " HeapAlloc:\t%d Bytes\n", s.HeapAlloc)
	var dummy string
	fmt.Printf("Enter to continue...")
	fmt.Scanf("%s", &dummy)
}
<commit_msg>cmd/trace: add basic documentation to main page<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"cmd/internal/browser"
	"flag"
	"fmt"
	"html/template"
	"internal/trace"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"runtime"
	"runtime/debug"
	"sync"

	_ "net/http/pprof" // Required to use pprof
)

const usageMessage = "" +
	`Usage of 'go tool trace':
Given a trace file produced by 'go test':
	go test -trace=trace.out pkg

Open a web browser displaying trace:
	go tool trace [flags] [pkg.test] trace.out

Generate a pprof-like profile from the trace:
    go tool trace -pprof=TYPE [pkg.test] trace.out

[pkg.test] argument is required for traces produced by Go 1.6 and below.
Go 1.7 does not require the binary argument.

Supported profile types are:
    - net: network blocking profile
    - sync: synchronization blocking profile
    - syscall: syscall blocking profile
    - sched: scheduler latency profile

Flags:
	-http=addr: HTTP service address (e.g., ':6060')
	-pprof=type: print a pprof-like profile instead
	-d: print debug info such as parsed events

Note that while the various profiles available when launching
'go tool trace' work on every browser, the trace viewer itself
(the 'view trace' page) comes from the Chrome/Chromium project
and is only actively tested on that browser.
`

var (
	httpFlag  = flag.String("http", "localhost:0", "HTTP service address (e.g., ':6060')")
	pprofFlag = flag.String("pprof", "", "print a pprof-like profile instead")
	debugFlag = flag.Bool("d", false, "print debug information such as parsed events list")

	// The binary file name, left here for serveSVGProfile.
	programBinary string
	traceFile     string
)

func main() {
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, usageMessage)
		os.Exit(2)
	}
	flag.Parse()

	// Go 1.7 traces embed symbol info and does not require the binary.
	// But we optionally accept binary as first arg for Go 1.5 traces.
	switch flag.NArg() {
	case 1:
		traceFile = flag.Arg(0)
	case 2:
		programBinary = flag.Arg(0)
		traceFile = flag.Arg(1)
	default:
		flag.Usage()
	}

	var pprofFunc func(io.Writer, *http.Request) error
	switch *pprofFlag {
	case "net":
		pprofFunc = pprofByGoroutine(computePprofIO)
	case "sync":
		pprofFunc = pprofByGoroutine(computePprofBlock)
	case "syscall":
		pprofFunc = pprofByGoroutine(computePprofSyscall)
	case "sched":
		pprofFunc = pprofByGoroutine(computePprofSched)
	}
	if pprofFunc != nil {
		if err := pprofFunc(os.Stdout, &http.Request{}); err != nil {
			dief("failed to generate pprof: %v\n", err)
		}
		os.Exit(0)
	}
	if *pprofFlag != "" {
		dief("unknown pprof type %s\n", *pprofFlag)
	}

	ln, err := net.Listen("tcp", *httpFlag)
	if err != nil {
		dief("failed to create server socket: %v\n", err)
	}

	log.Print("Parsing trace...")
	res, err := parseTrace()
	if err != nil {
		dief("%v\n", err)
	}

	if *debugFlag {
		trace.Print(res.Events)
		os.Exit(0)
	}
	reportMemoryUsage("after parsing trace")
	debug.FreeOSMemory()

	log.Print("Splitting trace...")
	ranges = splitTrace(res)
	reportMemoryUsage("after spliting trace")
	debug.FreeOSMemory()

	addr := "http://" + ln.Addr().String()
	log.Printf("Opening browser. Trace viewer is listening on %s", addr)
	browser.Open(addr)

	// Start http server.
	http.HandleFunc("/", httpMain)
	err = http.Serve(ln, nil)
	dief("failed to start http server: %v\n", err)
}

var ranges []Range

var loader struct {
	once sync.Once
	res  trace.ParseResult
	err  error
}

// parseEvents is a compatibility wrapper that returns only
// the Events part of trace.ParseResult returned by parseTrace.
func parseEvents() ([]*trace.Event, error) {
	res, err := parseTrace()
	if err != nil {
		return nil, err
	}
	return res.Events, err
}

func parseTrace() (trace.ParseResult, error) {
	loader.once.Do(func() {
		tracef, err := os.Open(traceFile)
		if err != nil {
			loader.err = fmt.Errorf("failed to open trace file: %v", err)
			return
		}
		defer tracef.Close()

		// Parse and symbolize.
		res, err := trace.Parse(bufio.NewReader(tracef), programBinary)
		if err != nil {
			loader.err = fmt.Errorf("failed to parse trace: %v", err)
			return
		}
		loader.res = res
	})
	return loader.res, loader.err
}

// httpMain serves the starting page.
func httpMain(w http.ResponseWriter, r *http.Request) {
	if err := templMain.Execute(w, ranges); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

var templMain = template.Must(template.New("").Parse(`

cmd/trace: the Go trace event viewer

  This web server provides various visualizations of an event log gathered during
  the execution of a Go program that uses the runtime/trace package.

Event timelines for running goroutines

{{if $}}

  Large traces are split into multiple sections of equal data size
  (not duration) to avoid overwhelming the visualizer.

	{{range $e := $}}

View trace ({{$e.Name}})

	{{end}}

{{else}}

View trace

{{end}}

  This view displays a timeline for each of the GOMAXPROCS logical
  processors, showing which goroutine (if any) was running on that
  logical processor at each moment.

  Each goroutine has an identifying number (e.g. G123), main function,
  and color.

  A colored bar represents an uninterrupted span of execution.

  Execution of a goroutine may migrate from one logical processor to another,
  causing a single colored bar to be horizontally continuous but
  vertically displaced.

  Clicking on a span reveals information about it, such as its
  duration, its causal predecessors and successors, and the stack trace
  at the final moment when it yielded the logical processor, for example
  because it made a system call or tried to acquire a mutex.

  Directly underneath each bar, a smaller bar or more commonly a fine
  vertical line indicates an event occuring during its execution.
  Some of these are related to garbage collection; most indicate that
  a goroutine yielded its logical processor but then immediately resumed execution
  on the same logical processor. Clicking on the event displays the stack trace
  at the moment it occurred.

  The causal relationships between spans of goroutine execution
  can be displayed by clicking the Flow Events button at the top.

  At the top ("STATS"), there are three additional timelines that
  display statistical information.

  "Goroutines" is a time series of the count of existing goroutines;
  clicking on it displays their breakdown by state at that moment:
  running, runnable, or waiting.

  "Heap" is a time series of the amount of heap memory allocated (in orange)
  and (in green) the allocation limit at which the next GC cycle will begin.

  "Threads" shows the number of kernel threads in existence: there is
  always one kernel thread per logical processor, and additional threads
  are created for calls to non-Go code such as a system call or a
  function written in C.

  Above the event trace for the first logical processor are
  traces for various runtime-internal events.

  The "GC" bar shows when the garbage collector is running, and in which stage.
  Garbage collection may temporarily affect all the logical processors
  and the other metrics.

  The "Network", "Timers", and "Syscalls" traces indicate events in
  the runtime that cause goroutines to wake up.

  The visualization allows you to navigate events at scales ranging from several
  seconds to a handful of nanoseconds.

  Consult the documentation for the Chromium Trace Event Profiling Tool
  for help navigating the view.

Goroutine analysis

  This view displays information about each set of goroutines that
  shares the same main function.

  Clicking on a main function shows links to the four types of
  blocking profile (see below) applied to that subset of goroutines.

  It also shows a table of specific goroutine instances, with various
  execution statistics and a link to the event timeline for each one.

  The timeline displays only the selected goroutine and any others it
  interacts with via block/unblock events. (The timeline is
  goroutine-oriented rather than logical processor-oriented.)

Profiles

  Each link below displays a global profile in zoomable graph form as
  produced by pprof's "web" command.

  In addition there is a link to download the profile for offline
  analysis with pprof.

  All four profiles represent causes of delay that prevent a goroutine
  from running on a logical processor: because it was waiting for the network,
  for a synchronization operation on a mutex or channel, for a system call,
  or for a logical processor to become available.

Network blocking profile (⬇)

Synchronization blocking profile (⬇)

Syscall blocking profile (⬇)

Scheduler latency profile (⬇)

User-defined tasks and regions

  The trace API allows a target program to annotate a region of code
  within a goroutine, such as a key function, so that its performance
  can be analyzed.

  Log events may be
  associated with a region to record progress and relevant values.

  The API also allows annotation of higher-level
  tasks,
  which may involve work across many goroutines.

  The links below display, for each region and task, a histogram of its execution times.

  Each histogram bucket contains a sample trace that records the
  sequence of events such as goroutine creations, log events, and
  subregion start/end times.

  For each task, you can click through to a logical-processor or
  goroutine-oriented view showing the tasks and regions on the
  timeline.

  Such information may help uncover which steps in a region are
  unexpectedly slow, or reveal relationships between the data values
  logged in a request and its running time.

User-defined tasks

User-defined regions

Garbage collection metrics

Minimum mutator utilization

  This chart indicates the maximum GC pause time (the largest x value
  for which y is zero), and more generally, the fraction of time that
  the processors are available to application goroutines ("mutators"),
  for any time window of a specified size, in the worst case.

`))

func dief(msg string, args ...any) {
	fmt.Fprintf(os.Stderr, msg, args...)
	os.Exit(1)
}

var debugMemoryUsage bool

func init() {
	v := os.Getenv("DEBUG_MEMORY_USAGE")
	debugMemoryUsage = v != ""
}

func reportMemoryUsage(msg string) {
	if !debugMemoryUsage {
		return
	}
	var s runtime.MemStats
	runtime.ReadMemStats(&s)
	w := os.Stderr
	fmt.Fprintf(w, "%s\n", msg)
	fmt.Fprintf(w, " Alloc:\t%d Bytes\n", s.Alloc)
	fmt.Fprintf(w, " Sys:\t%d Bytes\n", s.Sys)
	fmt.Fprintf(w, " HeapReleased:\t%d Bytes\n", s.HeapReleased)
	fmt.Fprintf(w, " HeapSys:\t%d Bytes\n", s.HeapSys)
	fmt.Fprintf(w, " HeapInUse:\t%d Bytes\n", s.HeapInuse)
	fmt.Fprintf(w, " HeapAlloc:\t%d Bytes\n", s.HeapAlloc)
	var dummy string
	fmt.Printf("Enter to continue...")
	fmt.Scanf("%s", &dummy)
}
<|endoftext|>


<commit_before>package receptor

import (
	"encoding/json"
	"github.com/blang/receptor/config"
	"github.com/blang/receptor/discovery"
	"github.com/blang/receptor/event"
	"github.com/blang/receptor/handler"
	"github.com/blang/receptor/reactor"
	"strconv"
	"strings"
	"testing"
	"time"
)

type testWatcher struct {
	setupCalled  bool
	acceptCalled bool
}

func (w *testWatcher) Setup(json.RawMessage) error {
	w.setupCalled = true
	return nil
}

func (w *testWatcher) Accept(json.RawMessage) (handler.Handler, error) {
	w.acceptCalled = true
	return handler.HandlerFunc(func(eventCh chan event.Event, closeCh chan struct{}) {
		for i := 0; i < 100; i++ {
			eventCh <- &event.SingleNode{
				EName: "test" + strconv.Itoa(i),
				EType: event.EventNodeUp,
			}
		}
	}), nil
}

type testReactor struct {
	setupCalled    bool
	acceptCalled   bool
	receivedEvents []event.Event
	eventRedirect  chan event.Event
}

func (r *testReactor) Setup(json.RawMessage) error {
	r.setupCalled = true
	return nil
}
func (r *testReactor) Accept(json.RawMessage) (handler.Handler, error) {
	r.acceptCalled = true
	return handler.HandlerFunc(func(eventCh chan event.Event, closeCh chan struct{}) {
		for e := range eventCh {
			r.receivedEvents = append(r.receivedEvents, e)
			r.eventRedirect <- e
		}
		close(r.eventRedirect)
	}), nil
}

func TestSystem(t *testing.T) {
	watcher := &testWatcher{}
	react := &testReactor{
		eventRedirect: make(chan event.Event),
	}
	discovery.Watchers["testWatcher"] = watcher
	reactor.Reactors["testReactor"] = react

	//config
	serviceConfig := config.ServiceConfig{
		Watchers: make(map[string]config.ActorConfig),
		Reactors: make(map[string]config.ActorConfig),
	}
	serviceConfig.Watchers["testWatcher1"] = config.ActorConfig{
		Type:   "testWatcher",
		Config: nil,
	}
	serviceConfig.Reactors["testReactor1"] = config.ActorConfig{
		Type:   "testReactor",
		Config: nil,
	}

	cfg := config.Config{
		Services: make(map[string]config.ServiceConfig),
		Watchers: make(map[string]json.RawMessage),
		Reactors: make(map[string]json.RawMessage),
	}

	cfg.Watchers["testWatcher"] = nil
	cfg.Reactors["testReactor"] = nil
	cfg.Services["testService"] = serviceConfig
	receptor := NewReceptor()
	err := receptor.Init(cfg)
	if err != nil {
		t.Fatalf("Init returned error: %s", err)
	}
	receptor.Start()

	count := 0
	timeout := time.After(5 * time.Second)
	for {
		select {
		case e, ok := <-react.eventRedirect:
			if !ok {
				return
			}
			count += len(e.Nodes())
			t.Logf("%d Event received: %s\n", len(e.Nodes()), stringEventNodes(e.Nodes()))
			if count == 100 {
				receptor.Stop()
			}
		case <-timeout:
			t.Error("Timeout")
			return
			// t.Fail("Timeout")

		}

	}

}

func stringEventNodes(nodes []event.NodeData) string {
	var str []string
	for _, n := range nodes {
		str = append(str, n.Name())
	}
	return strings.Join(str, ", ")
}
<commit_msg>Code cleanup<commit_after>package receptor

import (
	"encoding/json"
	"github.com/blang/receptor/config"
	"github.com/blang/receptor/discovery"
	"github.com/blang/receptor/event"
	"github.com/blang/receptor/handler"
	"github.com/blang/receptor/reactor"
	"strconv"
	"strings"
	"testing"
	"time"
)

type testWatcher struct {
	setupCalled  bool
	acceptCalled bool
}

func (w *testWatcher) Setup(json.RawMessage) error {
	w.setupCalled = true
	return nil
}

func (w *testWatcher) Accept(json.RawMessage) (handler.Handler, error) {
	w.acceptCalled = true
	return handler.HandlerFunc(func(eventCh chan event.Event, closeCh chan struct{}) {
		for i := 0; i < 100; i++ {
			eventCh <- &event.SingleNode{
				EName: "test" + strconv.Itoa(i),
				EType: event.EventNodeUp,
			}
		}
	}), nil
}

type testReactor struct {
	setupCalled    bool
	acceptCalled   bool
	receivedEvents []event.Event
	eventRedirect  chan event.Event
}

func (r *testReactor) Setup(json.RawMessage) error {
	r.setupCalled = true
	return nil
}
func (r *testReactor) Accept(json.RawMessage) (handler.Handler, error) {
	r.acceptCalled = true
	return handler.HandlerFunc(func(eventCh chan event.Event, closeCh chan struct{}) {
		for e := range eventCh {
			r.receivedEvents = append(r.receivedEvents, e)
			r.eventRedirect <- e
		}
		close(r.eventRedirect)
	}), nil
}

func TestSystem(t *testing.T) {
	watcher := &testWatcher{}
	react := &testReactor{
		eventRedirect: make(chan event.Event),
	}
	discovery.Watchers["testWatcher"] = watcher
	reactor.Reactors["testReactor"] = react

	//config
	serviceConfig := config.ServiceConfig{
		Watchers: make(map[string]config.ActorConfig),
		Reactors: make(map[string]config.ActorConfig),
	}
	serviceConfig.Watchers["testWatcher1"] = config.ActorConfig{
		Type:   "testWatcher",
		Config: nil,
	}
	serviceConfig.Reactors["testReactor1"] = config.ActorConfig{
		Type:   "testReactor",
		Config: nil,
	}

	cfg := config.Config{
		Services: make(map[string]config.ServiceConfig),
		Watchers: make(map[string]json.RawMessage),
		Reactors: make(map[string]json.RawMessage),
	}

	cfg.Watchers["testWatcher"] = nil
	cfg.Reactors["testReactor"] = nil
	cfg.Services["testService"] = serviceConfig
	receptor := NewReceptor()
	err := receptor.Init(cfg)
	if err != nil {
		t.Fatalf("Init returned error: %s", err)
	}
	receptor.Start()

	count := 0
	timeout := time.After(5 * time.Second)
	for {
		select {
		case e, ok := <-react.eventRedirect:
			if !ok {
				return
			}
			count += len(e.Nodes())
			t.Logf("%d Event received: %s\n", len(e.Nodes()), stringEventNodes(e.Nodes()))
			if count == 100 {
				receptor.Stop()
			}
		case <-timeout:
			t.Error("Timeout")
			return
		}
	}
}

func stringEventNodes(nodes []event.NodeData) string {
	var str []string
	for _, n := range nodes {
		str = append(str, n.Name())
	}
	return strings.Join(str, ", ")
}
<|endoftext|>


<commit_before>package redlot

import "github.com/syndtr/goleveldb/leveldb/util"

func encodeHashKey(name, key []byte) (buf []byte) {
	buf = append(buf, typeHASH)
	buf = append(buf, uint32ToBytes(uint32(len(name)))...)
	buf = append(buf, name...)
	buf = append(buf, uint32ToBytes(uint32(len(key)))...)
	buf = append(buf, key...)
	return
}

func decodeHashKey(b []byte) (name, key []byte) {
	nameLen := bytesToUint32(b[1:5])
	name = b[5 : 5+nameLen]
	key = b[9+nameLen:]
	return
}

func encodeHsizeKey(name []byte) (buf []byte) {
	buf = append(buf, typeHSIZE)
	buf = append(buf, uint32ToBytes(uint32(len(name)))...)
	buf = append(buf, name...)
	return
}

func decodeHsizeKey(b []byte) (key []byte) {
	return b[5:]
}

func hashSizeIncr(name, key []byte) {
	hsize := encodeHsizeKey(name)
	hash := encodeHashKey(name, key)

	var size uint32
	if b, err := db.Get(hsize, nil); err == nil {
		size = bytesToUint32(b)
	}

	if exists, _ := db.Has(hash, nil); !exists {
		size++
		db.Put(hsize, uint32ToBytes(size), nil)
	}
}

// Hset will set a hashmap value by the key.
// Args: name string, key string, value any
func Hset(args [][]byte) (r interface{}, err error) {
	if len(args) < 3 {
		return nil, errNosArgs
	}

	hashSizeIncr(args[0], args[1])
	err = db.Put(encodeHashKey(args[0], args[1]), args[2], nil)
	if err != nil {
		return nil, err
	}

	return
}

// Hset will return a hashmap value by the key.
// Args: name string, key string
func Hget(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}
	var b []byte
	b, err = db.Get(encodeHashKey(args[0], args[1]), nil)
	return string(b), err
}

// Hdel will delete a hashmap value by the key.
// Args: name string, key string
func Hdel(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}

// Hincr will incr a hashmap value by the key.
// Args: name string, key string, value int
func Hincr(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}

// Hexists will check the hashmap key is exists.
// Args: name string, key string
func Hexists(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}

// Hsize will return the hashmap size.
// Args: name string
func Hsize(args [][]byte) (r interface{}, err error) {
	if len(args) < 1 {
		return nil, errNosArgs
	}

	var size int64
	if b, e := db.Get(encodeHsizeKey(args[0]), nil); e != nil {
		size = -1
	} else {
		size = int64(bytesToUint32(b))
	}

	return size, nil
}

// Hlist will list all hashmap in the range.
// Args: start string, end string, limit int
func Hlist(args [][]byte) (r []string, err error) {
	if len(args) < 3 {
		return nil, errNosArgs
	}

	return
}

// Hrlist will reverse list all hashmap in the range.
// Args: start string, end string, limit int
func Hrlist(args [][]byte) (r []string, err error) {
	if len(args) < 3 {
		return nil, errNosArgs
	}

	return
}

// Hkeys will list the hashmap keys in the range.
// Args: name string, start string, end string, limit int
func Hkeys(args [][]byte) (r []string, err error) {
	if len(args) < 4 {
		return nil, errNosArgs
	}

	return
}

// Hgetall will list all keys/value in the hashmap.
// Args: name string
func Hgetall(args [][]byte) (r []string, err error) {
	if len(args) < 1 {
		return nil, errNosArgs
	}

	if _, err = db.Get(encodeHsizeKey(args[0]), nil); err != nil {
		return
	}

	var buf []byte
	buf = append(buf, typeHASH)
	buf = append(buf, uint32ToBytes(uint32(len(args[0])))...)
	buf = append(buf, args[0]...)
	ke := append(buf, []byte{0xff}...)

	iter := db.NewIterator(&util.Range{Start: buf, Limit: ke}, nil)
	for iter.Next() {
		_, key := decodeHashKey(iter.Key())
		r = append(r, string(key))
		r = append(r, string(iter.Value()))
	}
	iter.Release()
	err = iter.Error()
	return
}

// Hscan will list keys/value of the hashmap in the range.
// Args: name string, start string, end string, limit int
func Hscan(args [][]byte) (r []string, err error) {
	if len(args) < 4 {
		return nil, errNosArgs
	}

	return
}

// Hrscan will reverse list keys/value of the hashmap in the range.
// Args: name string, start string, end string, limit int
func Hrscan(args [][]byte) (r []string, err error) {
	if len(args) < 4 {
		return nil, errNosArgs
	}

	return
}

// Hclear will remove all value in the hashmap.
// Args: name string
func Hclear(args [][]byte) (r interface{}, err error) {
	if len(args) < 1 {
		return nil, errNosArgs
	}

	return
}

// MultiHget will return multi hashmap value by keys.
func MultiHget(args [][]byte) (r []string, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}

// MultiHset will set multi hashmap value by keys.
func MultiHset(args [][]byte) (r interface{}, err error) {
	if len(args) < 3 && len(args)%2 == 0 {
		return nil, errNosArgs
	}

	return
}

// MultiHdel will delete multi hashmap value by keys.
func MultiHdel(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}
<commit_msg>Add methods to increment hash value.<commit_after>package redlot

import (
	"strconv"

	"github.com/syndtr/goleveldb/leveldb/util"
)

func encodeHashKey(name, key []byte) (buf []byte) {
	buf = append(buf, typeHASH)
	buf = append(buf, uint32ToBytes(uint32(len(name)))...)
	buf = append(buf, name...)
	buf = append(buf, uint32ToBytes(uint32(len(key)))...)
	buf = append(buf, key...)
	return
}

func decodeHashKey(b []byte) (name, key []byte) {
	nameLen := bytesToUint32(b[1:5])
	name = b[5 : 5+nameLen]
	key = b[9+nameLen:]
	return
}

func encodeHsizeKey(name []byte) (buf []byte) {
	buf = append(buf, typeHSIZE)
	buf = append(buf, uint32ToBytes(uint32(len(name)))...)
	buf = append(buf, name...)
	return
}

func decodeHsizeKey(b []byte) (key []byte) {
	return b[5:]
}

func hashSizeIncr(name, key []byte) {
	hsize := encodeHsizeKey(name)
	hash := encodeHashKey(name, key)

	var size uint32
	if b, err := db.Get(hsize, nil); err == nil {
		size = bytesToUint32(b)
	}

	if exists, _ := db.Has(hash, nil); !exists {
		size++
		db.Put(hsize, uint32ToBytes(size), nil)
	}
}

// Hset will set a hashmap value by the key.
// Args: name string, key string, value any
func Hset(args [][]byte) (r interface{}, err error) {
	if len(args) < 3 {
		return nil, errNosArgs
	}

	hashSizeIncr(args[0], args[1])
	err = db.Put(encodeHashKey(args[0], args[1]), args[2], nil)
	if err != nil {
		return nil, err
	}

	return
}

// Hset will return a hashmap value by the key.
// Args: name string, key string
func Hget(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}
	var b []byte
	b, err = db.Get(encodeHashKey(args[0], args[1]), nil)
	return string(b), err
}

// Hdel will delete a hashmap value by the key.
// Args: name string, key string
func Hdel(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}

func hincr(key []byte, increment int) (r int64, err error) {
	v, _ := db.Get(key, nil)
	var number int
	if len(v) != 0 {
		var err error
		number, err = strconv.Atoi(string(v))
		if err != nil {
			return -1, errNotInt
		}
	}
	number += increment
	return int64(number), db.Put(key, []byte(strconv.Itoa(number)), nil)
}

// Hincr will incr a hashmap value by the key.
// Args: name string, key string
func Hincr(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}
	key := encodeHashKey(args[0], args[1])

	return hincr(key, 1)
}
}

// Hexists will check the hashmap key is exists.
// Args: name string, key string
func Hexists(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}

// Hsize will return the hashmap size.
// Args: name string
func Hsize(args [][]byte) (r interface{}, err error) {
	if len(args) < 1 {
		return nil, errNosArgs
	}

	var size int64
	if b, e := db.Get(encodeHsizeKey(args[0]), nil); e != nil {
		size = -1
	} else {
		size = int64(bytesToUint32(b))
	}

	return size, nil
}

// Hlist will list all hashmap in the range.
// Args: start string, end string, limit int
func Hlist(args [][]byte) (r []string, err error) {
	if len(args) < 3 {
		return nil, errNosArgs
	}

	return
}

// Hrlist will reverse list all hashmap in the range.
// Args: start string, end string, limit int
func Hrlist(args [][]byte) (r []string, err error) {
	if len(args) < 3 {
		return nil, errNosArgs
	}

	return
}

// Hkeys will list the hashmap keys in the range.
// Args: name string, start string, end string, limit int
func Hkeys(args [][]byte) (r []string, err error) {
	if len(args) < 4 {
		return nil, errNosArgs
	}

	return
}

// Hgetall will list all keys/value in the hashmap.
// Args: name string
func Hgetall(args [][]byte) (r []string, err error) {
	if len(args) < 1 {
		return nil, errNosArgs
	}

	if _, err = db.Get(encodeHsizeKey(args[0]), nil); err != nil {
		return
	}

	var buf []byte
	buf = append(buf, typeHASH)
	buf = append(buf, uint32ToBytes(uint32(len(args[0])))...)
	buf = append(buf, args[0]...)
	ke := append(buf, []byte{0xff}...)

	iter := db.NewIterator(&util.Range{Start: buf, Limit: ke}, nil)
	for iter.Next() {
		_, key := decodeHashKey(iter.Key())
		r = append(r, string(key))
		r = append(r, string(iter.Value()))
	}
	iter.Release()
	err = iter.Error()
	return
}

// Hscan will list keys/value of the hashmap in the range.
// Args: name string, start string, end string, limit int
func Hscan(args [][]byte) (r []string, err error) {
	if len(args) < 4 {
		return nil, errNosArgs
	}

	return
}

// Hrscan will reverse list keys/value of the hashmap in the range.
// Args: name string, start string, end string, limit int
func Hrscan(args [][]byte) (r []string, err error) {
	if len(args) < 4 {
		return nil, errNosArgs
	}

	return
}

// Hclear will remove all value in the hashmap.
// Args: name string
func Hclear(args [][]byte) (r interface{}, err error) {
	if len(args) < 1 {
		return nil, errNosArgs
	}

	return
}

// MultiHget will return multi hashmap value by keys.
func MultiHget(args [][]byte) (r []string, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}

// MultiHset will set multi hashmap value by keys.
func MultiHset(args [][]byte) (r interface{}, err error) {
	if len(args) < 3 && len(args)%2 == 0 {
		return nil, errNosArgs
	}

	return
}

// MultiHdel will delete multi hashmap value by keys.
func MultiHdel(args [][]byte) (r interface{}, err error) {
	if len(args) < 2 {
		return nil, errNosArgs
	}

	return
}
<|endoftext|>


<commit_before>// Copyright (C) 2015  The GoHBase Authors.  All rights reserved.
// This file is part of GoHBase.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package region

import (
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/golang/protobuf/proto"
	"github.com/tsuna/gohbase/hrpc"
	"github.com/tsuna/gohbase/pb"
)

// ClientType is a type alias to represent the type of this region client
type ClientType string

var (
	// ErrShortWrite is used when the writer thread only succeeds in writing
	// part of its buffer to the socket, and not all of the buffer was sent
	ErrShortWrite = errors.New("short write occurred while writing to socket")

	// ErrMissingCallID is used when HBase sends us a response message for a
	// request that we didn't send
	ErrMissingCallID = errors.New("HBase responded to a nonsensical call ID")

	// javaRetryableExceptions is a map where all Java exceptions that signify
	// the RPC should be sent again are listed (as keys). If a Java exception
	// listed here is returned by HBase, the client should attempt to resend
	// the RPC message, potentially via a different region client.
	javaRetryableExceptions = map[string]struct{}{
		"org.apache.hadoop.hbase.NotServingRegionException":         struct{}{},
		"org.apache.hadoop.hbase.exceptions.RegionMovedException":   struct{}{},
		"org.apache.hadoop.hbase.exceptions.RegionOpeningException": struct{}{},
	}
)

const (
	// RegionClient is a ClientType that means this will be a normal client
	RegionClient = ClientType("ClientService")

	// MasterClient is a ClientType that means this client will talk to the
	// master server
	MasterClient = ClientType("MasterService")
)

// UnrecoverableError is an error that this region.Client can't recover from.
// The connection to the RegionServer has to be closed and all queued and
// outstanding RPCs will be failed / retried.
type UnrecoverableError struct {
	error
}

func (e UnrecoverableError) Error() string {
	return e.error.Error()
}

// RetryableError is an error that indicates the RPC should be retried because
// the error is transient (e.g. a region being momentarily unavailable).
type RetryableError struct {
	error
}

func (e RetryableError) Error() string {
	return e.error.Error()
}

// Client manages a connection to a RegionServer.
type Client struct {
	id uint32

	conn net.Conn

	// Hostname or IP address of the RegionServer.
	host string

	// Port of the RegionServer.
	port uint16

	// writeMutex is used to prevent multiple threads from writing to the
	// socket at the same time.
	writeMutex *sync.Mutex

	// sendErr is set once a write fails.
	sendErr     error
	sendErrLock sync.Mutex

	rpcs []hrpc.Call

	// Once the rpcs list has grown to a large enough size, this channel is
	// written to to notify the writer thread that it should stop sleeping and
	// process the list
	process chan struct{}

	// sentRPCs contains the mapping of sent call IDs to RPC calls, so that when
	// a response is received it can be tied to the correct RPC
	sentRPCs      map[uint32]hrpc.Call
	sentRPCsMutex *sync.Mutex

	rpcQueueSize  int
	flushInterval time.Duration
}

// NewClient creates a new RegionClient.
func NewClient(host string, port uint16, ctype ClientType,
	queueSize int, flushInterval time.Duration) (*Client, error) {
	addr := fmt.Sprintf("%s:%d", host, port)
	conn, err := net.Dial("tcp", addr) // TODO: DialTimeout
	if err != nil {
		return nil,
			fmt.Errorf("failed to connect to the RegionServer at %s: %s", addr, err)
	}
	c := &Client{
		conn:          conn,
		host:          host,
		port:          port,
		writeMutex:    &sync.Mutex{},
		process:       make(chan struct{}),
		sentRPCsMutex: &sync.Mutex{},
		sentRPCs:      make(map[uint32]hrpc.Call),
		rpcQueueSize:  queueSize,
		flushInterval: flushInterval,
	}
	err = c.sendHello(ctype)
	if err != nil {
		return nil, err
	}
	go c.processRpcs() // Writer goroutine
	go c.receiveRpcs() // Reader goroutine
	return c, nil
}

// Close asks this region.Client to close its connection to the RegionServer.
// All queued and outstanding RPCs, if any, will be failed as if a connection
// error had happened.
func (c *Client) Close() {
	c.setSendErr(errors.New("shutting down"))
	c.errorEncountered()
}

// Host returns the host that this client talks to
func (c *Client) Host() string {
	return c.host
}

// Port returns the port that this client talks over
func (c *Client) Port() uint16 {
	return c.port
}

func (c *Client) getSendErr() error {
	c.sendErrLock.Lock()
	err := c.sendErr
	c.sendErrLock.Unlock()
	return err
}

func (c *Client) setSendErr(err error) {
	c.sendErrLock.Lock()
	c.sendErr = err
	c.sendErrLock.Unlock()
}

func (c *Client) processRpcs() {
	for {
		if c.getSendErr() != nil {
			return
		}

		select {
		case <-time.After(c.flushInterval):
			select {
			case <-c.process:
			// If we got a message on c.process at the same time as our
			// timeout elapsed, we'll non-deterministically land in either
			// cases of this outer select.  Here we double-check whether
			// something was written onto c.process, in which case we don't
			// grab the lock (see comment below in the other case).
			default:
				c.writeMutex.Lock()
			}
		case <-c.process:
			// We don't acquire the lock here, because the thread that sent
			// something on the process channel will have locked the mutex,
			// and will not release it so as to transfer ownership
		}

		rpcs := make([]hrpc.Call, len(c.rpcs))
		for i, rpc := range c.rpcs {
			rpcs[i] = rpc
		}
		c.rpcs = nil
		c.writeMutex.Unlock()

		for i, rpc := range rpcs {
			// If the deadline has been exceeded, don't bother sending the
			// request. The function that placed the RPC in our queue should
			// stop waiting for a result and return an error.
			select {
			case _, ok := <-rpc.GetContext().Done():
				if !ok {
					continue
				}
			default:
			}

			err := c.sendRPC(rpc)
			if err != nil {
				_, ok := err.(UnrecoverableError)
				if ok {
					c.setSendErr(err)

					c.writeMutex.Lock()
					c.rpcs = append(c.rpcs, rpcs[i:]...)
					c.writeMutex.Unlock()

					c.errorEncountered()
					return
				}
				rpc.GetResultChan() <- hrpc.RPCResult{Error: err}
			}
		}
	}
}

func (c *Client) receiveRpcs() {
	var sz [4]byte
	for {
		err := c.readFully(sz[:])
		if err != nil {
			c.setSendErr(err)
			c.errorEncountered()
			return
		}

		buf := make([]byte, binary.BigEndian.Uint32(sz[:]))
		err = c.readFully(buf)
		if err != nil {
			c.setSendErr(err)
			c.errorEncountered()
			return
		}

		resp := &pb.ResponseHeader{}
		respLen, nb := proto.DecodeVarint(buf)
		buf = buf[nb:]
		err = proto.UnmarshalMerge(buf[:respLen], resp)
		buf = buf[respLen:]
		if err != nil {
			// Failed to deserialize the response header
			c.setSendErr(err)
			c.errorEncountered()
			return
		}
		if resp.CallId == nil {
			// Response doesn't have a call ID
			log.Error("Response doesn't have a call ID!")
			c.setSendErr(ErrMissingCallID)
			c.errorEncountered()
			return
		}

		c.sentRPCsMutex.Lock()
		rpc, ok := c.sentRPCs[*resp.CallId]
		c.sentRPCsMutex.Unlock()

		if !ok {
			log.WithFields(log.Fields{
				"CallId": *resp.CallId,
			}).Error("Received a response with an unexpected call ID")

			log.Error("Waiting for responses to the following calls:")
			c.sentRPCsMutex.Lock()
			for id, call := range c.sentRPCs {
				log.Errorf("\t\t%d: %v", id, call)
			}
			c.sentRPCsMutex.Unlock()

			c.setSendErr(fmt.Errorf("HBase sent a response with an unexpected call ID: %d",
				resp.CallId))
			c.errorEncountered()
			return
		}

		var rpcResp proto.Message
		if resp.Exception == nil {
			respLen, nb = proto.DecodeVarint(buf)
			buf = buf[nb:]
			rpcResp = rpc.NewResponse()
			err = proto.UnmarshalMerge(buf, rpcResp)
			buf = buf[respLen:]
		} else {
			javaClass := *resp.Exception.ExceptionClassName
			err = fmt.Errorf("HBase Java exception %s: \n%s", javaClass,
				*resp.Exception.StackTrace)
			if _, ok := javaRetryableExceptions[javaClass]; ok {
				// This is a recoverable error. The client should retry.
				err = RetryableError{err}
			}
		}
		rpc.GetResultChan() <- hrpc.RPCResult{Msg: rpcResp, Error: err}

		c.sentRPCsMutex.Lock()
		delete(c.sentRPCs, *resp.CallId)
		c.sentRPCsMutex.Unlock()
	}
}

func (c *Client) errorEncountered() {
	c.writeMutex.Lock()
	res := hrpc.RPCResult{Error: UnrecoverableError{c.getSendErr()}}
	for _, rpc := range c.rpcs {
		rpc.GetResultChan() <- res
	}
	c.rpcs = nil
	c.writeMutex.Unlock()

	c.sentRPCsMutex.Lock()
	for _, rpc := range c.sentRPCs {
		rpc.GetResultChan() <- res
	}
	c.sentRPCs = nil
	c.sentRPCsMutex.Unlock()

	c.conn.Close()
}

// Sends the given buffer to the RegionServer.
func (c *Client) write(buf []byte) error {
	n, err := c.conn.Write(buf)

	if err != nil {
		// There was an error while writing
		return err
	}
	if n != len(buf) {
		// We failed to write the entire buffer
		// TODO: Perhaps handle this in another way than closing down
		return ErrShortWrite
	}
	return nil
}

// Tries to read enough data to fully fill up the given buffer.
func (c *Client) readFully(buf []byte) error {
	// TODO: Handle short reads.
	n, err := c.conn.Read(buf)
	if err != nil {
		return fmt.Errorf("Failed to read from the RS: %s", err)
	} else if n != len(buf) {
		return fmt.Errorf("Failed to read everything from the RS: %s", err)
	}
	return nil
}

// Sends the "hello" message needed when opening a new connection.
func (c *Client) sendHello(ctype ClientType) error {
	connHeader := &pb.ConnectionHeader{
		UserInfo: &pb.UserInformation{
			EffectiveUser: proto.String("gopher"),
		},
		ServiceName: proto.String(string(ctype)),
		//CellBlockCodecClass: "org.apache.hadoop.hbase.codec.KeyValueCodec",
	}
	data, err := proto.Marshal(connHeader)
	if err != nil {
		return fmt.Errorf("failed to marshal connection header: %s", err)
	}

	const header = "HBas\x00\x50" // \x50 = Simple Auth.
	buf := make([]byte, 0, len(header)+4+len(data))
	buf = append(buf, header...)
	buf = buf[:len(header)+4]
	binary.BigEndian.PutUint32(buf[6:], uint32(len(data)))
	buf = append(buf, data...)

	return c.write(buf)
}

// QueueRPC will add an rpc call to the queue for processing by the writer
// goroutine
func (c *Client) QueueRPC(rpc hrpc.Call) error {
	sendErr := c.getSendErr()
	if sendErr != nil {
		return sendErr
	}
	c.writeMutex.Lock()
	c.rpcs = append(c.rpcs, rpc)
	if len(c.rpcs) > c.rpcQueueSize {
		c.process <- struct{}{}
		// We don't release the lock here, because we want to transfer ownership
		// of the lock to the goroutine that processes the RPCs
	} else {
		c.writeMutex.Unlock()
	}
	return nil
}

// sendRPC sends an RPC out to the wire.
// Returns the response (for now, as the call is synchronous).
func (c *Client) sendRPC(rpc hrpc.Call) error {
	// Header.
	c.id++
	reqheader := &pb.RequestHeader{
		CallId:       &c.id,
		MethodName:   proto.String(rpc.GetName()),
		RequestParam: proto.Bool(true),
	}

	payload, err := rpc.Serialize()
	if err != nil {
		return fmt.Errorf("Failed to serialize RPC: %s", err)
	}
	payloadLen := proto.EncodeVarint(uint64(len(payload)))

	headerData, err := proto.Marshal(reqheader)
	if err != nil {
		return fmt.Errorf("Failed to marshal Get request: %s", err)
	}

	buf := make([]byte, 5, 4+1+len(headerData)+len(payloadLen)+len(payload))
	binary.BigEndian.PutUint32(buf, uint32(cap(buf)-4))
	buf[4] = byte(len(headerData))
	buf = append(buf, headerData...)
	buf = append(buf, payloadLen...)
	buf = append(buf, payload...)

	c.sentRPCsMutex.Lock()
	c.sentRPCs[c.id] = rpc
	c.sentRPCsMutex.Unlock()

	err = c.write(buf)
	if err != nil {
		return UnrecoverableError{err}
	}

	return nil
}
<commit_msg>region: readFully to actually read fully<commit_after>// Copyright (C) 2015  The GoHBase Authors.  All rights reserved.
// This file is part of GoHBase.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package region

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/golang/protobuf/proto"
	"github.com/tsuna/gohbase/hrpc"
	"github.com/tsuna/gohbase/pb"
)

// ClientType is a type alias to represent the type of this region client
type ClientType string

var (
	// ErrShortWrite is used when the writer thread only succeeds in writing
	// part of its buffer to the socket, and not all of the buffer was sent
	ErrShortWrite = errors.New("short write occurred while writing to socket")

	// ErrMissingCallID is used when HBase sends us a response message for a
	// request that we didn't send
	ErrMissingCallID = errors.New("HBase responded to a nonsensical call ID")

	// javaRetryableExceptions is a map where all Java exceptions that signify
	// the RPC should be sent again are listed (as keys). If a Java exception
	// listed here is returned by HBase, the client should attempt to resend
	// the RPC message, potentially via a different region client.
	javaRetryableExceptions = map[string]struct{}{
		"org.apache.hadoop.hbase.NotServingRegionException":         struct{}{},
		"org.apache.hadoop.hbase.exceptions.RegionMovedException":   struct{}{},
		"org.apache.hadoop.hbase.exceptions.RegionOpeningException": struct{}{},
	}
)

const (
	// RegionClient is a ClientType that means this will be a normal client
	RegionClient = ClientType("ClientService")

	// MasterClient is a ClientType that means this client will talk to the
	// master server
	MasterClient = ClientType("MasterService")
)

// UnrecoverableError is an error that this region.Client can't recover from.
// The connection to the RegionServer has to be closed and all queued and
// outstanding RPCs will be failed / retried.
type UnrecoverableError struct {
	error
}

func (e UnrecoverableError) Error() string {
	return e.error.Error()
}

// RetryableError is an error that indicates the RPC should be retried because
// the error is transient (e.g. a region being momentarily unavailable).
type RetryableError struct {
	error
}

func (e RetryableError) Error() string {
	return e.error.Error()
}

// Client manages a connection to a RegionServer.
type Client struct {
	id uint32

	conn net.Conn

	// Hostname or IP address of the RegionServer.
	host string

	// Port of the RegionServer.
	port uint16

	// writeMutex is used to prevent multiple threads from writing to the
	// socket at the same time.
	writeMutex *sync.Mutex

	// sendErr is set once a write fails.
	sendErr     error
	sendErrLock sync.Mutex

	rpcs []hrpc.Call

	// Once the rpcs list has grown to a large enough size, this channel is
	// written to to notify the writer thread that it should stop sleeping and
	// process the list
	process chan struct{}

	// sentRPCs contains the mapping of sent call IDs to RPC calls, so that when
	// a response is received it can be tied to the correct RPC
	sentRPCs      map[uint32]hrpc.Call
	sentRPCsMutex *sync.Mutex

	rpcQueueSize  int
	flushInterval time.Duration
}

// NewClient creates a new RegionClient.
func NewClient(host string, port uint16, ctype ClientType,
	queueSize int, flushInterval time.Duration) (*Client, error) {
	addr := fmt.Sprintf("%s:%d", host, port)
	conn, err := net.Dial("tcp", addr) // TODO: DialTimeout
	if err != nil {
		return nil,
			fmt.Errorf("failed to connect to the RegionServer at %s: %s", addr, err)
	}
	c := &Client{
		conn:          conn,
		host:          host,
		port:          port,
		writeMutex:    &sync.Mutex{},
		process:       make(chan struct{}),
		sentRPCsMutex: &sync.Mutex{},
		sentRPCs:      make(map[uint32]hrpc.Call),
		rpcQueueSize:  queueSize,
		flushInterval: flushInterval,
	}
	err = c.sendHello(ctype)
	if err != nil {
		return nil, err
	}
	go c.processRpcs() // Writer goroutine
	go c.receiveRpcs() // Reader goroutine
	return c, nil
}

// Close asks this region.Client to close its connection to the RegionServer.
// All queued and outstanding RPCs, if any, will be failed as if a connection
// error had happened.
func (c *Client) Close() {
	c.setSendErr(errors.New("shutting down"))
	c.errorEncountered()
}

// Host returns the host that this client talks to
func (c *Client) Host() string {
	return c.host
}

// Port returns the port that this client talks over
func (c *Client) Port() uint16 {
	return c.port
}

func (c *Client) getSendErr() error {
	c.sendErrLock.Lock()
	err := c.sendErr
	c.sendErrLock.Unlock()
	return err
}

func (c *Client) setSendErr(err error) {
	c.sendErrLock.Lock()
	c.sendErr = err
	c.sendErrLock.Unlock()
}

func (c *Client) processRpcs() {
	for {
		if c.getSendErr() != nil {
			return
		}

		select {
		case <-time.After(c.flushInterval):
			select {
			case <-c.process:
			// If we got a message on c.process at the same time as our
			// timeout elapsed, we'll non-deterministically land in either
			// cases of this outer select.  Here we double-check whether
			// something was written onto c.process, in which case we don't
			// grab the lock (see comment below in the other case).
			default:
				c.writeMutex.Lock()
			}
		case <-c.process:
			// We don't acquire the lock here, because the thread that sent
			// something on the process channel will have locked the mutex,
			// and will not release it so as to transfer ownership
		}

		rpcs := make([]hrpc.Call, len(c.rpcs))
		for i, rpc := range c.rpcs {
			rpcs[i] = rpc
		}
		c.rpcs = nil
		c.writeMutex.Unlock()

		for i, rpc := range rpcs {
			// If the deadline has been exceeded, don't bother sending the
			// request. The function that placed the RPC in our queue should
			// stop waiting for a result and return an error.
			select {
			case _, ok := <-rpc.GetContext().Done():
				if !ok {
					continue
				}
			default:
			}

			err := c.sendRPC(rpc)
			if err != nil {
				_, ok := err.(UnrecoverableError)
				if ok {
					c.setSendErr(err)

					c.writeMutex.Lock()
					c.rpcs = append(c.rpcs, rpcs[i:]...)
					c.writeMutex.Unlock()

					c.errorEncountered()
					return
				}
				rpc.GetResultChan() <- hrpc.RPCResult{Error: err}
			}
		}
	}
}

func (c *Client) receiveRpcs() {
	var sz [4]byte
	for {
		err := c.readFully(sz[:])
		if err != nil {
			c.setSendErr(err)
			c.errorEncountered()
			return
		}

		buf := make([]byte, binary.BigEndian.Uint32(sz[:]))
		err = c.readFully(buf)
		if err != nil {
			c.setSendErr(err)
			c.errorEncountered()
			return
		}

		resp := &pb.ResponseHeader{}
		respLen, nb := proto.DecodeVarint(buf)
		buf = buf[nb:]
		err = proto.UnmarshalMerge(buf[:respLen], resp)
		buf = buf[respLen:]
		if err != nil {
			// Failed to deserialize the response header
			c.setSendErr(err)
			c.errorEncountered()
			return
		}
		if resp.CallId == nil {
			// Response doesn't have a call ID
			log.Error("Response doesn't have a call ID!")
			c.setSendErr(ErrMissingCallID)
			c.errorEncountered()
			return
		}

		c.sentRPCsMutex.Lock()
		rpc, ok := c.sentRPCs[*resp.CallId]
		c.sentRPCsMutex.Unlock()

		if !ok {
			log.WithFields(log.Fields{
				"CallId": *resp.CallId,
			}).Error("Received a response with an unexpected call ID")

			log.Error("Waiting for responses to the following calls:")
			c.sentRPCsMutex.Lock()
			for id, call := range c.sentRPCs {
				log.Errorf("\t\t%d: %v", id, call)
			}
			c.sentRPCsMutex.Unlock()

			c.setSendErr(fmt.Errorf("HBase sent a response with an unexpected call ID: %d",
				resp.CallId))
			c.errorEncountered()
			return
		}

		var rpcResp proto.Message
		if resp.Exception == nil {
			respLen, nb = proto.DecodeVarint(buf)
			buf = buf[nb:]
			rpcResp = rpc.NewResponse()
			err = proto.UnmarshalMerge(buf, rpcResp)
			buf = buf[respLen:]
		} else {
			javaClass := *resp.Exception.ExceptionClassName
			err = fmt.Errorf("HBase Java exception %s: \n%s", javaClass,
				*resp.Exception.StackTrace)
			if _, ok := javaRetryableExceptions[javaClass]; ok {
				// This is a recoverable error. The client should retry.
				err = RetryableError{err}
			}
		}
		rpc.GetResultChan() <- hrpc.RPCResult{Msg: rpcResp, Error: err}

		c.sentRPCsMutex.Lock()
		delete(c.sentRPCs, *resp.CallId)
		c.sentRPCsMutex.Unlock()
	}
}

func (c *Client) errorEncountered() {
	c.writeMutex.Lock()
	res := hrpc.RPCResult{Error: UnrecoverableError{c.getSendErr()}}
	for _, rpc := range c.rpcs {
		rpc.GetResultChan() <- res
	}
	c.rpcs = nil
	c.writeMutex.Unlock()

	c.sentRPCsMutex.Lock()
	for _, rpc := range c.sentRPCs {
		rpc.GetResultChan() <- res
	}
	c.sentRPCs = nil
	c.sentRPCsMutex.Unlock()

	c.conn.Close()
}

// Sends the given buffer to the RegionServer.
func (c *Client) write(buf []byte) error {
	n, err := c.conn.Write(buf)

	if err != nil {
		// There was an error while writing
		return err
	}
	if n != len(buf) {
		// We failed to write the entire buffer
		// TODO: Perhaps handle this in another way than closing down
		return ErrShortWrite
	}
	return nil
}

// Tries to read enough data to fully fill up the given buffer.
func (c *Client) readFully(buf []byte) error {
	_, err := io.ReadFull(c.conn, buf)
	if err != nil {
		return fmt.Errorf("Failed to read from the RS: %s", err)
	}
	return nil
}

// Sends the "hello" message needed when opening a new connection.
func (c *Client) sendHello(ctype ClientType) error {
	connHeader := &pb.ConnectionHeader{
		UserInfo: &pb.UserInformation{
			EffectiveUser: proto.String("gopher"),
		},
		ServiceName: proto.String(string(ctype)),
		//CellBlockCodecClass: "org.apache.hadoop.hbase.codec.KeyValueCodec",
	}
	data, err := proto.Marshal(connHeader)
	if err != nil {
		return fmt.Errorf("failed to marshal connection header: %s", err)
	}

	const header = "HBas\x00\x50" // \x50 = Simple Auth.
	buf := make([]byte, 0, len(header)+4+len(data))
	buf = append(buf, header...)
	buf = buf[:len(header)+4]
	binary.BigEndian.PutUint32(buf[6:], uint32(len(data)))
	buf = append(buf, data...)

	return c.write(buf)
}

// QueueRPC will add an rpc call to the queue for processing by the writer
// goroutine
func (c *Client) QueueRPC(rpc hrpc.Call) error {
	sendErr := c.getSendErr()
	if sendErr != nil {
		return sendErr
	}
	c.writeMutex.Lock()
	c.rpcs = append(c.rpcs, rpc)
	if len(c.rpcs) > c.rpcQueueSize {
		c.process <- struct{}{}
		// We don't release the lock here, because we want to transfer ownership
		// of the lock to the goroutine that processes the RPCs
	} else {
		c.writeMutex.Unlock()
	}
	return nil
}

// sendRPC sends an RPC out to the wire.
// Returns the response (for now, as the call is synchronous).
func (c *Client) sendRPC(rpc hrpc.Call) error {
	// Header.
	c.id++
	reqheader := &pb.RequestHeader{
		CallId:       &c.id,
		MethodName:   proto.String(rpc.GetName()),
		RequestParam: proto.Bool(true),
	}

	payload, err := rpc.Serialize()
	if err != nil {
		return fmt.Errorf("Failed to serialize RPC: %s", err)
	}
	payloadLen := proto.EncodeVarint(uint64(len(payload)))

	headerData, err := proto.Marshal(reqheader)
	if err != nil {
		return fmt.Errorf("Failed to marshal Get request: %s", err)
	}

	buf := make([]byte, 5, 4+1+len(headerData)+len(payloadLen)+len(payload))
	binary.BigEndian.PutUint32(buf, uint32(cap(buf)-4))
	buf[4] = byte(len(headerData))
	buf = append(buf, headerData...)
	buf = append(buf, payloadLen...)
	buf = append(buf, payload...)

	c.sentRPCsMutex.Lock()
	c.sentRPCs[c.id] = rpc
	c.sentRPCsMutex.Unlock()

	err = c.write(buf)
	if err != nil {
		return UnrecoverableError{err}
	}

	return nil
}
<|endoftext|>


<commit_before>package registry

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/cliconfig"
)

// Login tries to register/login to the registry server.
func Login(authConfig *cliconfig.AuthConfig, registryEndpoint *Endpoint) (string, error) {
	// Separates the v2 registry login logic from the v1 logic.
	if registryEndpoint.Version == APIVersion2 {
		return loginV2(authConfig, registryEndpoint, "" /* scope */)
	}
	return loginV1(authConfig, registryEndpoint)
}

// loginV1 tries to register/login to the v1 registry server.
func loginV1(authConfig *cliconfig.AuthConfig, registryEndpoint *Endpoint) (string, error) {
	var (
		status        string
		reqBody       []byte
		err           error
		reqStatusCode = 0
		serverAddress = authConfig.ServerAddress
	)

	logrus.Debugf("attempting v1 login to registry endpoint %s", registryEndpoint)

	if serverAddress == "" {
		return "", fmt.Errorf("Server Error: Server Address not set.")
	}

	loginAgainstOfficialIndex := serverAddress == IndexServer

	// to avoid sending the server address to the server it should be removed before being marshalled
	authCopy := *authConfig
	authCopy.ServerAddress = ""

	jsonBody, err := json.Marshal(authCopy)
	if err != nil {
		return "", fmt.Errorf("Config Error: %s", err)
	}

	// using `bytes.NewReader(jsonBody)` here causes the server to respond with a 411 status.
	b := strings.NewReader(string(jsonBody))
	req1, err := registryEndpoint.client.Post(serverAddress+"users/", "application/json; charset=utf-8", b)
	if err != nil {
		return "", fmt.Errorf("Server Error: %s", err)
	}
	defer req1.Body.Close()
	reqStatusCode = req1.StatusCode
	reqBody, err = ioutil.ReadAll(req1.Body)
	if err != nil {
		return "", fmt.Errorf("Server Error: [%#v] %s", reqStatusCode, err)
	}

	if reqStatusCode == 201 {
		if loginAgainstOfficialIndex {
			status = "Account created. Please use the confirmation link we sent" +
				" to your e-mail to activate it."
		} else {
			// *TODO: Use registry configuration to determine what this says, if anything?
			status = "Account created. Please see the documentation of the registry " + serverAddress + " for instructions how to activate it."
		}
	} else if reqStatusCode == 400 {
		if string(reqBody) == "\"Username or email already exists\"" {
			req, err := http.NewRequest("GET", serverAddress+"users/", nil)
			req.SetBasicAuth(authConfig.Username, authConfig.Password)
			resp, err := registryEndpoint.client.Do(req)
			if err != nil {
				return "", err
			}
			defer resp.Body.Close()
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return "", err
			}
			if resp.StatusCode == 200 {
				return "Login Succeeded", nil
			} else if resp.StatusCode == 401 {
				return "", fmt.Errorf("Wrong login/password, please try again")
			} else if resp.StatusCode == 403 {
				if loginAgainstOfficialIndex {
					return "", fmt.Errorf("Login: Account is not Active. Please check your e-mail for a confirmation link.")
				}
				// *TODO: Use registry configuration to determine what this says, if anything?
				return "", fmt.Errorf("Login: Account is not Active. Please see the documentation of the registry %s for instructions how to activate it.", serverAddress)
			} else if resp.StatusCode == 500 { // Issue #14326
				logrus.Errorf("%s returned status code %d. Response Body :\n%s", req.URL.String(), resp.StatusCode, body)
				return "", fmt.Errorf("Internal Server Error")
			}
			return "", fmt.Errorf("Login: %s (Code: %d; Headers: %s)", body, resp.StatusCode, resp.Header)
		}
		return "", fmt.Errorf("Registration: %s", reqBody)

	} else if reqStatusCode == 401 {
		// This case would happen with private registries where /v1/users is
		// protected, so people can use `docker login` as an auth check.
		req, err := http.NewRequest("GET", serverAddress+"users/", nil)
		req.SetBasicAuth(authConfig.Username, authConfig.Password)
		resp, err := registryEndpoint.client.Do(req)
		if err != nil {
			return "", err
		}
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return "", err
		}
		if resp.StatusCode == 200 {
			return "Login Succeeded", nil
		} else if resp.StatusCode == 401 {
			return "", fmt.Errorf("Wrong login/password, please try again")
		} else {
			return "", fmt.Errorf("Login: %s (Code: %d; Headers: %s)", body,
				resp.StatusCode, resp.Header)
		}
	} else {
		return "", fmt.Errorf("Unexpected status code [%d] : %s", reqStatusCode, reqBody)
	}
	return status, nil
}

// loginV2 tries to login to the v2 registry server. The given registry endpoint has been
// pinged or setup with a list of authorization challenges. Each of these challenges are
// tried until one of them succeeds. Currently supported challenge schemes are:
// 		HTTP Basic Authorization
// 		Token Authorization with a separate token issuing server
// NOTE: the v2 logic does not attempt to create a user account if one doesn't exist. For
// now, users should create their account through other means like directly from a web page
// served by the v2 registry service provider. Whether this will be supported in the future
// is to be determined.
func loginV2(authConfig *cliconfig.AuthConfig, registryEndpoint *Endpoint, scope string) (string, error) {
	logrus.Debugf("attempting v2 login to registry endpoint %s", registryEndpoint)
	var (
		err       error
		allErrors []error
	)

	for _, challenge := range registryEndpoint.AuthChallenges {
		params := make(map[string]string, len(challenge.Parameters)+1)
		for k, v := range challenge.Parameters {
			params[k] = v
		}
		params["scope"] = scope
		logrus.Debugf("trying %q auth challenge with params %v", challenge.Scheme, params)

		switch strings.ToLower(challenge.Scheme) {
		case "basic":
			err = tryV2BasicAuthLogin(authConfig, params, registryEndpoint)
		case "bearer":
			err = tryV2TokenAuthLogin(authConfig, params, registryEndpoint)
		default:
			// Unsupported challenge types are explicitly skipped.
			err = fmt.Errorf("unsupported auth scheme: %q", challenge.Scheme)
		}

		if err == nil {
			return "Login Succeeded", nil
		}

		logrus.Debugf("error trying auth challenge %q: %s", challenge.Scheme, err)

		allErrors = append(allErrors, err)
	}

	return "", fmt.Errorf("no successful auth challenge for %s - errors: %s", registryEndpoint, allErrors)
}

func tryV2BasicAuthLogin(authConfig *cliconfig.AuthConfig, params map[string]string, registryEndpoint *Endpoint) error {
	req, err := http.NewRequest("GET", registryEndpoint.Path(""), nil)
	if err != nil {
		return err
	}

	req.SetBasicAuth(authConfig.Username, authConfig.Password)

	resp, err := registryEndpoint.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("basic auth attempt to %s realm %q failed with status: %d %s", registryEndpoint, params["realm"], resp.StatusCode, http.StatusText(resp.StatusCode))
	}

	return nil
}

func tryV2TokenAuthLogin(authConfig *cliconfig.AuthConfig, params map[string]string, registryEndpoint *Endpoint) error {
	token, err := getToken(authConfig.Username, authConfig.Password, params, registryEndpoint)
	if err != nil {
		return err
	}

	req, err := http.NewRequest("GET", registryEndpoint.Path(""), nil)
	if err != nil {
		return err
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))

	resp, err := registryEndpoint.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("token auth attempt to %s realm %q failed with status: %d %s", registryEndpoint, params["realm"], resp.StatusCode, http.StatusText(resp.StatusCode))
	}

	return nil
}

// ResolveAuthConfig matches an auth configuration to a server address or a URL
func ResolveAuthConfig(config *cliconfig.ConfigFile, index *IndexInfo) cliconfig.AuthConfig {
	configKey := index.GetAuthConfigKey()
	// First try the happy case
	if c, found := config.AuthConfigs[configKey]; found || index.Official {
		return c
	}

	convertToHostname := func(url string) string {
		stripped := url
		if strings.HasPrefix(url, "http://") {
			stripped = strings.Replace(url, "http://", "", 1)
		} else if strings.HasPrefix(url, "https://") {
			stripped = strings.Replace(url, "https://", "", 1)
		}

		nameParts := strings.SplitN(stripped, "/", 2)

		return nameParts[0]
	}

	// Maybe they have a legacy config file, we will iterate the keys converting
	// them to the new format and testing
	for registry, ac := range config.AuthConfigs {
		if configKey == convertToHostname(registry) {
			return ac
		}
	}

	// When all else fails, return an empty auth config
	return cliconfig.AuthConfig{}
}
<commit_msg>rename req to resp<commit_after>package registry

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/cliconfig"
)

// Login tries to register/login to the registry server.
func Login(authConfig *cliconfig.AuthConfig, registryEndpoint *Endpoint) (string, error) {
	// Separates the v2 registry login logic from the v1 logic.
	if registryEndpoint.Version == APIVersion2 {
		return loginV2(authConfig, registryEndpoint, "" /* scope */)
	}
	return loginV1(authConfig, registryEndpoint)
}

// loginV1 tries to register/login to the v1 registry server.
func loginV1(authConfig *cliconfig.AuthConfig, registryEndpoint *Endpoint) (string, error) {
	var (
		status         string
		respBody       []byte
		err            error
		respStatusCode = 0
		serverAddress  = authConfig.ServerAddress
	)

	logrus.Debugf("attempting v1 login to registry endpoint %s", registryEndpoint)

	if serverAddress == "" {
		return "", fmt.Errorf("Server Error: Server Address not set.")
	}

	loginAgainstOfficialIndex := serverAddress == IndexServer

	// to avoid sending the server address to the server it should be removed before being marshalled
	authCopy := *authConfig
	authCopy.ServerAddress = ""

	jsonBody, err := json.Marshal(authCopy)
	if err != nil {
		return "", fmt.Errorf("Config Error: %s", err)
	}

	// using `bytes.NewReader(jsonBody)` here causes the server to respond with a 411 status.
	b := strings.NewReader(string(jsonBody))
	resp1, err := registryEndpoint.client.Post(serverAddress+"users/", "application/json; charset=utf-8", b)
	if err != nil {
		return "", fmt.Errorf("Server Error: %s", err)
	}
	defer resp1.Body.Close()
	respStatusCode = resp1.StatusCode
	respBody, err = ioutil.ReadAll(resp1.Body)
	if err != nil {
		return "", fmt.Errorf("Server Error: [%#v] %s", respStatusCode, err)
	}

	if respStatusCode == 201 {
		if loginAgainstOfficialIndex {
			status = "Account created. Please use the confirmation link we sent" +
				" to your e-mail to activate it."
		} else {
			// *TODO: Use registry configuration to determine what this says, if anything?
			status = "Account created. Please see the documentation of the registry " + serverAddress + " for instructions how to activate it."
		}
	} else if respStatusCode == 400 {
		if string(respBody) == "\"Username or email already exists\"" {
			req, err := http.NewRequest("GET", serverAddress+"users/", nil)
			req.SetBasicAuth(authConfig.Username, authConfig.Password)
			resp, err := registryEndpoint.client.Do(req)
			if err != nil {
				return "", err
			}
			defer resp.Body.Close()
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return "", err
			}
			if resp.StatusCode == 200 {
				return "Login Succeeded", nil
			} else if resp.StatusCode == 401 {
				return "", fmt.Errorf("Wrong login/password, please try again")
			} else if resp.StatusCode == 403 {
				if loginAgainstOfficialIndex {
					return "", fmt.Errorf("Login: Account is not Active. Please check your e-mail for a confirmation link.")
				}
				// *TODO: Use registry configuration to determine what this says, if anything?
				return "", fmt.Errorf("Login: Account is not Active. Please see the documentation of the registry %s for instructions how to activate it.", serverAddress)
			} else if resp.StatusCode == 500 { // Issue #14326
				logrus.Errorf("%s returned status code %d. Response Body :\n%s", req.URL.String(), resp.StatusCode, body)
				return "", fmt.Errorf("Internal Server Error")
			}
			return "", fmt.Errorf("Login: %s (Code: %d; Headers: %s)", body, resp.StatusCode, resp.Header)
		}
		return "", fmt.Errorf("Registration: %s", respBody)

	} else if respStatusCode == 401 {
		// This case would happen with private registries where /v1/users is
		// protected, so people can use `docker login` as an auth check.
		req, err := http.NewRequest("GET", serverAddress+"users/", nil)
		req.SetBasicAuth(authConfig.Username, authConfig.Password)
		resp, err := registryEndpoint.client.Do(req)
		if err != nil {
			return "", err
		}
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return "", err
		}
		if resp.StatusCode == 200 {
			return "Login Succeeded", nil
		} else if resp.StatusCode == 401 {
			return "", fmt.Errorf("Wrong login/password, please try again")
		} else {
			return "", fmt.Errorf("Login: %s (Code: %d; Headers: %s)", body,
				resp.StatusCode, resp.Header)
		}
	} else {
		return "", fmt.Errorf("Unexpected status code [%d] : %s", respStatusCode, respBody)
	}
	return status, nil
}

// loginV2 tries to login to the v2 registry server. The given registry endpoint has been
// pinged or setup with a list of authorization challenges. Each of these challenges are
// tried until one of them succeeds. Currently supported challenge schemes are:
// 		HTTP Basic Authorization
// 		Token Authorization with a separate token issuing server
// NOTE: the v2 logic does not attempt to create a user account if one doesn't exist. For
// now, users should create their account through other means like directly from a web page
// served by the v2 registry service provider. Whether this will be supported in the future
// is to be determined.
func loginV2(authConfig *cliconfig.AuthConfig, registryEndpoint *Endpoint, scope string) (string, error) {
	logrus.Debugf("attempting v2 login to registry endpoint %s", registryEndpoint)
	var (
		err       error
		allErrors []error
	)

	for _, challenge := range registryEndpoint.AuthChallenges {
		params := make(map[string]string, len(challenge.Parameters)+1)
		for k, v := range challenge.Parameters {
			params[k] = v
		}
		params["scope"] = scope
		logrus.Debugf("trying %q auth challenge with params %v", challenge.Scheme, params)

		switch strings.ToLower(challenge.Scheme) {
		case "basic":
			err = tryV2BasicAuthLogin(authConfig, params, registryEndpoint)
		case "bearer":
			err = tryV2TokenAuthLogin(authConfig, params, registryEndpoint)
		default:
			// Unsupported challenge types are explicitly skipped.
			err = fmt.Errorf("unsupported auth scheme: %q", challenge.Scheme)
		}

		if err == nil {
			return "Login Succeeded", nil
		}

		logrus.Debugf("error trying auth challenge %q: %s", challenge.Scheme, err)

		allErrors = append(allErrors, err)
	}

	return "", fmt.Errorf("no successful auth challenge for %s - errors: %s", registryEndpoint, allErrors)
}

func tryV2BasicAuthLogin(authConfig *cliconfig.AuthConfig, params map[string]string, registryEndpoint *Endpoint) error {
	req, err := http.NewRequest("GET", registryEndpoint.Path(""), nil)
	if err != nil {
		return err
	}

	req.SetBasicAuth(authConfig.Username, authConfig.Password)

	resp, err := registryEndpoint.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("basic auth attempt to %s realm %q failed with status: %d %s", registryEndpoint, params["realm"], resp.StatusCode, http.StatusText(resp.StatusCode))
	}

	return nil
}

func tryV2TokenAuthLogin(authConfig *cliconfig.AuthConfig, params map[string]string, registryEndpoint *Endpoint) error {
	token, err := getToken(authConfig.Username, authConfig.Password, params, registryEndpoint)
	if err != nil {
		return err
	}

	req, err := http.NewRequest("GET", registryEndpoint.Path(""), nil)
	if err != nil {
		return err
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))

	resp, err := registryEndpoint.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("token auth attempt to %s realm %q failed with status: %d %s", registryEndpoint, params["realm"], resp.StatusCode, http.StatusText(resp.StatusCode))
	}

	return nil
}

// ResolveAuthConfig matches an auth configuration to a server address or a URL
func ResolveAuthConfig(config *cliconfig.ConfigFile, index *IndexInfo) cliconfig.AuthConfig {
	configKey := index.GetAuthConfigKey()
	// First try the happy case
	if c, found := config.AuthConfigs[configKey]; found || index.Official {
		return c
	}

	convertToHostname := func(url string) string {
		stripped := url
		if strings.HasPrefix(url, "http://") {
			stripped = strings.Replace(url, "http://", "", 1)
		} else if strings.HasPrefix(url, "https://") {
			stripped = strings.Replace(url, "https://", "", 1)
		}

		nameParts := strings.SplitN(stripped, "/", 2)

		return nameParts[0]
	}

	// Maybe they have a legacy config file, we will iterate the keys converting
	// them to the new format and testing
	for registry, ac := range config.AuthConfigs {
		if configKey == convertToHostname(registry) {
			return ac
		}
	}

	// When all else fails, return an empty auth config
	return cliconfig.AuthConfig{}
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
	"time"
)

func listIndexFn(t *testing.T, ifn IndexFn, f File, fs chan File, es chan error) (files []File) {
	go func() {
		ifn(f, fs, es)
	}()
	for {
		select {
		case f := <-fs:
			files = append(files, f)
		case e := <-es:
			if e != nil {
				t.Fatal(e)
			}
		case <-time.After(1 * time.Second):
			return
		}
	}
}
<commit_msg>add generic testing for index functions<commit_after>package main

import (
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func testIndexFn(t *testing.T, ifn IndexFn, f File, expected []File) {
	fs := make(chan File)
	es := make(chan error)
	actual := []File{}
	go func() {
		ifn(f, fs, es)
	}()
	for {
		select {
		case f := <-fs:
			actual = append(actual, f)
		case e := <-es:
			if e != nil {
				t.Error(e)
			}
		case <-time.After(1 * time.Second):
			t.Fatal("Index function timed out")
		}
	}
	for i, ac := range actual {
		exp := expected[i]

		expCont, e := exp.ReadAll()
		if e != nil {
			t.Fatal("Reading Expected File failed: ", e)
		}
		acContent, e := ac.ReadAll()
		if e != nil {
			t.Errorf("Reading %v failed: %v", ac, e)
		}

		switch {
		case string(expCont) != string(acContent):
			t.Errorf("%s != %s", expCont, acContent)
			fallthrough
		case !(ac.Mtime.Equal(exp.Mtime) && ac.Url == exp.Url):
			t.Errorf("%v != %v", exp, ac)
		}
	}
}

type testCase struct {
	in  File
	exp []File
}

func testIndex(t *testing.T, urlVar *string, ifn IndexFn,
	tests []testCase, tFn func(*testing.T, []File) http.Handler) {
	for _, test := range tests {
		s := httptest.NewServer(tFn(t, test.exp))
		urlVar = &s.URL
		testIndexFn(t, ifn, test.in, test.exp)
		s.Close()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
	"unsafe"

	"github.com/sirupsen/logrus"
	"github.com/spf13/pflag"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/debug"
	"golang.org/x/sys/windows/svc/eventlog"
	"golang.org/x/sys/windows/svc/mgr"
)

var (
	flServiceName       *string
	flRegisterService   *bool
	flUnregisterService *bool
	flRunService        *bool

	setStdHandle = windows.NewLazySystemDLL("kernel32.dll").NewProc("SetStdHandle")
	oldStderr    windows.Handle
	panicFile    *os.File

	service *handler
)

const (
	// These should match the values in event_messages.mc.
	eventInfo  = 1
	eventWarn  = 1
	eventError = 1
	eventDebug = 2
	eventPanic = 3
	eventFatal = 4

	eventExtraOffset = 10 // Add this to any event to get a string that supports extended data
)

func installServiceFlags(flags *pflag.FlagSet) {
	flServiceName = flags.String("service-name", "docker", "Set the Windows service name")
	flRegisterService = flags.Bool("register-service", false, "Register the service and exit")
	flUnregisterService = flags.Bool("unregister-service", false, "Unregister the service and exit")
	flRunService = flags.Bool("run-service", false, "")
	_ = flags.MarkHidden("run-service")
}

type handler struct {
	tosvc     chan bool
	fromsvc   chan error
	daemonCli *DaemonCli
}

type etwHook struct {
	log *eventlog.Log
}

func (h *etwHook) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.PanicLevel,
		logrus.FatalLevel,
		logrus.ErrorLevel,
		logrus.WarnLevel,
		logrus.InfoLevel,
		logrus.DebugLevel,
	}
}

func (h *etwHook) Fire(e *logrus.Entry) error {
	var (
		etype uint16
		eid   uint32
	)

	switch e.Level {
	case logrus.PanicLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventPanic
	case logrus.FatalLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventFatal
	case logrus.ErrorLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventError
	case logrus.WarnLevel:
		etype = windows.EVENTLOG_WARNING_TYPE
		eid = eventWarn
	case logrus.InfoLevel:
		etype = windows.EVENTLOG_INFORMATION_TYPE
		eid = eventInfo
	case logrus.DebugLevel:
		etype = windows.EVENTLOG_INFORMATION_TYPE
		eid = eventDebug
	default:
		return errors.New("unknown level")
	}

	// If there is additional data, include it as a second string.
	exts := ""
	if len(e.Data) > 0 {
		fs := bytes.Buffer{}
		for k, v := range e.Data {
			fs.WriteString(k)
			fs.WriteByte('=')
			fmt.Fprint(&fs, v)
			fs.WriteByte(' ')
		}

		exts = fs.String()[:fs.Len()-1]
		eid += eventExtraOffset
	}

	if h.log == nil {
		fmt.Fprintf(os.Stderr, "%s [%s]\n", e.Message, exts)
		return nil
	}

	var (
		ss  [2]*uint16
		err error
	)

	ss[0], err = windows.UTF16PtrFromString(e.Message)
	if err != nil {
		return err
	}

	count := uint16(1)
	if exts != "" {
		ss[1], err = windows.UTF16PtrFromString(exts)
		if err != nil {
			return err
		}

		count++
	}

	return windows.ReportEvent(h.log.Handle, etype, 0, eid, 0, count, 0, &ss[0], nil)
}

func getServicePath() (string, error) {
	p, err := exec.LookPath(os.Args[0])
	if err != nil {
		return "", err
	}
	return filepath.Abs(p)
}

func registerService() error {
	p, err := getServicePath()
	if err != nil {
		return err
	}
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	c := mgr.Config{
		ServiceType:  windows.SERVICE_WIN32_OWN_PROCESS,
		StartType:    mgr.StartAutomatic,
		ErrorControl: mgr.ErrorNormal,
		Dependencies: []string{},
		DisplayName:  "Docker Engine",
	}

	// Configure the service to launch with the arguments that were just passed.
	args := []string{"--run-service"}
	for _, a := range os.Args[1:] {
		if a != "--register-service" && a != "--unregister-service" {
			args = append(args, a)
		}
	}

	s, err := m.CreateService(*flServiceName, p, c, args...)
	if err != nil {
		return err
	}
	defer s.Close()

	// See http://stackoverflow.com/questions/35151052/how-do-i-configure-failure-actions-of-a-windows-service-written-in-go
	const (
		scActionNone       = 0
		scActionRestart    = 1
		scActionReboot     = 2
		scActionRunCommand = 3

		serviceConfigFailureActions = 2
	)

	type serviceFailureActions struct {
		ResetPeriod  uint32
		RebootMsg    *uint16
		Command      *uint16
		ActionsCount uint32
		Actions      uintptr
	}

	type scAction struct {
		Type  uint32
		Delay uint32
	}
	t := []scAction{
		{Type: scActionRestart, Delay: uint32(60 * time.Second / time.Millisecond)},
		{Type: scActionRestart, Delay: uint32(60 * time.Second / time.Millisecond)},
		{Type: scActionNone},
	}
	lpInfo := serviceFailureActions{ResetPeriod: uint32(24 * time.Hour / time.Second), ActionsCount: uint32(3), Actions: uintptr(unsafe.Pointer(&t[0]))}
	err = windows.ChangeServiceConfig2(s.Handle, serviceConfigFailureActions, (*byte)(unsafe.Pointer(&lpInfo)))
	if err != nil {
		return err
	}

	return eventlog.Install(*flServiceName, p, false, eventlog.Info|eventlog.Warning|eventlog.Error)
}

func unregisterService() error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	s, err := m.OpenService(*flServiceName)
	if err != nil {
		return err
	}
	defer s.Close()

	eventlog.Remove(*flServiceName)
	err = s.Delete()
	if err != nil {
		return err
	}
	return nil
}

// initService is the entry point for running the daemon as a Windows
// service. It returns an indication to stop (if registering/un-registering);
// an indication of whether it is running as a service; and an error.
func initService(daemonCli *DaemonCli) (bool, bool, error) {
	if *flUnregisterService {
		if *flRegisterService {
			return true, false, errors.New("--register-service and --unregister-service cannot be used together")
		}
		return true, false, unregisterService()
	}

	if *flRegisterService {
		return true, false, registerService()
	}

	if !*flRunService {
		return false, false, nil
	}

	interactive, err := svc.IsAnInteractiveSession()
	if err != nil {
		return false, false, err
	}

	h := &handler{
		tosvc:     make(chan bool),
		fromsvc:   make(chan error),
		daemonCli: daemonCli,
	}

	var log *eventlog.Log
	if !interactive {
		log, err = eventlog.Open(*flServiceName)
		if err != nil {
			return false, false, err
		}
	}

	logrus.AddHook(&etwHook{log})
	logrus.SetOutput(io.Discard)

	service = h
	go func() {
		if interactive {
			err = debug.Run(*flServiceName, h)
		} else {
			err = svc.Run(*flServiceName, h)
		}

		h.fromsvc <- err
	}()

	// Wait for the first signal from the service handler.
	err = <-h.fromsvc
	if err != nil {
		return false, false, err
	}
	return false, true, nil
}

func (h *handler) started() error {
	// This must be delayed until daemonCli initializes Config.Root
	err := initPanicFile(filepath.Join(h.daemonCli.Config.Root, "panic.log"))
	if err != nil {
		return err
	}

	h.tosvc <- false
	return nil
}

func (h *handler) stopped(err error) {
	logrus.Debugf("Stopping service: %v", err)
	h.tosvc <- err != nil
	<-h.fromsvc
}

func (h *handler) Execute(_ []string, r <-chan svc.ChangeRequest, s chan<- svc.Status) (bool, uint32) {
	s <- svc.Status{State: svc.StartPending, Accepts: 0}
	// Unblock initService()
	h.fromsvc <- nil

	// Wait for initialization to complete.
	failed := <-h.tosvc
	if failed {
		logrus.Debug("Aborting service start due to failure during initialization")
		return true, 1
	}

	s <- svc.Status{State: svc.Running, Accepts: svc.AcceptStop | svc.AcceptShutdown | svc.Accepted(windows.SERVICE_ACCEPT_PARAMCHANGE)}
	logrus.Debug("Service running")
Loop:
	for {
		select {
		case failed = <-h.tosvc:
			break Loop
		case c := <-r:
			switch c.Cmd {
			case svc.Cmd(windows.SERVICE_CONTROL_PARAMCHANGE):
				h.daemonCli.reloadConfig()
			case svc.Interrogate:
				s <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s <- svc.Status{State: svc.StopPending, Accepts: 0}
				h.daemonCli.stop()
			}
		}
	}

	removePanicFile()
	if failed {
		return true, 1
	}
	return false, 0
}

func initPanicFile(path string) error {
	var err error
	panicFile, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o200)
	if err != nil {
		return err
	}

	st, err := panicFile.Stat()
	if err != nil {
		return err
	}

	// If there are contents in the file already, move the file out of the way
	// and replace it.
	if st.Size() > 0 {
		panicFile.Close()
		os.Rename(path, path+".old")
		panicFile, err = os.Create(path)
		if err != nil {
			return err
		}
	}

	// Update STD_ERROR_HANDLE to point to the panic file so that Go writes to
	// it when it panics. Remember the old stderr to restore it before removing
	// the panic file.
	sh := uint32(windows.STD_ERROR_HANDLE)
	h, err := windows.GetStdHandle(sh)
	if err != nil {
		return err
	}

	oldStderr = h

	r, _, err := setStdHandle.Call(uintptr(sh), uintptr(panicFile.Fd()))
	if r == 0 && err != nil {
		return err
	}

	// Reset os.Stderr to the panic file (so fmt.Fprintf(os.Stderr,...) actually gets redirected)
	os.Stderr = os.NewFile(uintptr(panicFile.Fd()), "/dev/stderr")

	// Force threads that panic to write to stderr (the panicFile handle now), otherwise it will go into the ether
	log.SetOutput(os.Stderr)

	return nil
}

func removePanicFile() {
	if st, err := panicFile.Stat(); err == nil {
		if st.Size() == 0 {
			sh := uint32(windows.STD_ERROR_HANDLE)
			setStdHandle.Call(uintptr(sh), uintptr(oldStderr))
			panicFile.Close()
			os.Remove(panicFile.Name())
		}
	}
}
<commit_msg>cmd/dockerd: use golang.org/x/sys/windows.SetStdHandle()<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
	"unsafe"

	"github.com/sirupsen/logrus"
	"github.com/spf13/pflag"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/debug"
	"golang.org/x/sys/windows/svc/eventlog"
	"golang.org/x/sys/windows/svc/mgr"
)

var (
	flServiceName       *string
	flRegisterService   *bool
	flUnregisterService *bool
	flRunService        *bool

	oldStderr windows.Handle
	panicFile *os.File

	service *handler
)

const (
	// These should match the values in event_messages.mc.
	eventInfo  = 1
	eventWarn  = 1
	eventError = 1
	eventDebug = 2
	eventPanic = 3
	eventFatal = 4

	eventExtraOffset = 10 // Add this to any event to get a string that supports extended data
)

func installServiceFlags(flags *pflag.FlagSet) {
	flServiceName = flags.String("service-name", "docker", "Set the Windows service name")
	flRegisterService = flags.Bool("register-service", false, "Register the service and exit")
	flUnregisterService = flags.Bool("unregister-service", false, "Unregister the service and exit")
	flRunService = flags.Bool("run-service", false, "")
	_ = flags.MarkHidden("run-service")
}

type handler struct {
	tosvc     chan bool
	fromsvc   chan error
	daemonCli *DaemonCli
}

type etwHook struct {
	log *eventlog.Log
}

func (h *etwHook) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.PanicLevel,
		logrus.FatalLevel,
		logrus.ErrorLevel,
		logrus.WarnLevel,
		logrus.InfoLevel,
		logrus.DebugLevel,
	}
}

func (h *etwHook) Fire(e *logrus.Entry) error {
	var (
		etype uint16
		eid   uint32
	)

	switch e.Level {
	case logrus.PanicLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventPanic
	case logrus.FatalLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventFatal
	case logrus.ErrorLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventError
	case logrus.WarnLevel:
		etype = windows.EVENTLOG_WARNING_TYPE
		eid = eventWarn
	case logrus.InfoLevel:
		etype = windows.EVENTLOG_INFORMATION_TYPE
		eid = eventInfo
	case logrus.DebugLevel:
		etype = windows.EVENTLOG_INFORMATION_TYPE
		eid = eventDebug
	default:
		return errors.New("unknown level")
	}

	// If there is additional data, include it as a second string.
	exts := ""
	if len(e.Data) > 0 {
		fs := bytes.Buffer{}
		for k, v := range e.Data {
			fs.WriteString(k)
			fs.WriteByte('=')
			fmt.Fprint(&fs, v)
			fs.WriteByte(' ')
		}

		exts = fs.String()[:fs.Len()-1]
		eid += eventExtraOffset
	}

	if h.log == nil {
		fmt.Fprintf(os.Stderr, "%s [%s]\n", e.Message, exts)
		return nil
	}

	var (
		ss  [2]*uint16
		err error
	)

	ss[0], err = windows.UTF16PtrFromString(e.Message)
	if err != nil {
		return err
	}

	count := uint16(1)
	if exts != "" {
		ss[1], err = windows.UTF16PtrFromString(exts)
		if err != nil {
			return err
		}

		count++
	}

	return windows.ReportEvent(h.log.Handle, etype, 0, eid, 0, count, 0, &ss[0], nil)
}

func getServicePath() (string, error) {
	p, err := exec.LookPath(os.Args[0])
	if err != nil {
		return "", err
	}
	return filepath.Abs(p)
}

func registerService() error {
	p, err := getServicePath()
	if err != nil {
		return err
	}
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	c := mgr.Config{
		ServiceType:  windows.SERVICE_WIN32_OWN_PROCESS,
		StartType:    mgr.StartAutomatic,
		ErrorControl: mgr.ErrorNormal,
		Dependencies: []string{},
		DisplayName:  "Docker Engine",
	}

	// Configure the service to launch with the arguments that were just passed.
	args := []string{"--run-service"}
	for _, a := range os.Args[1:] {
		if a != "--register-service" && a != "--unregister-service" {
			args = append(args, a)
		}
	}

	s, err := m.CreateService(*flServiceName, p, c, args...)
	if err != nil {
		return err
	}
	defer s.Close()

	// See http://stackoverflow.com/questions/35151052/how-do-i-configure-failure-actions-of-a-windows-service-written-in-go
	const (
		scActionNone       = 0
		scActionRestart    = 1
		scActionReboot     = 2
		scActionRunCommand = 3

		serviceConfigFailureActions = 2
	)

	type serviceFailureActions struct {
		ResetPeriod  uint32
		RebootMsg    *uint16
		Command      *uint16
		ActionsCount uint32
		Actions      uintptr
	}

	type scAction struct {
		Type  uint32
		Delay uint32
	}
	t := []scAction{
		{Type: scActionRestart, Delay: uint32(60 * time.Second / time.Millisecond)},
		{Type: scActionRestart, Delay: uint32(60 * time.Second / time.Millisecond)},
		{Type: scActionNone},
	}
	lpInfo := serviceFailureActions{ResetPeriod: uint32(24 * time.Hour / time.Second), ActionsCount: uint32(3), Actions: uintptr(unsafe.Pointer(&t[0]))}
	err = windows.ChangeServiceConfig2(s.Handle, serviceConfigFailureActions, (*byte)(unsafe.Pointer(&lpInfo)))
	if err != nil {
		return err
	}

	return eventlog.Install(*flServiceName, p, false, eventlog.Info|eventlog.Warning|eventlog.Error)
}

func unregisterService() error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	s, err := m.OpenService(*flServiceName)
	if err != nil {
		return err
	}
	defer s.Close()

	eventlog.Remove(*flServiceName)
	err = s.Delete()
	if err != nil {
		return err
	}
	return nil
}

// initService is the entry point for running the daemon as a Windows
// service. It returns an indication to stop (if registering/un-registering);
// an indication of whether it is running as a service; and an error.
func initService(daemonCli *DaemonCli) (bool, bool, error) {
	if *flUnregisterService {
		if *flRegisterService {
			return true, false, errors.New("--register-service and --unregister-service cannot be used together")
		}
		return true, false, unregisterService()
	}

	if *flRegisterService {
		return true, false, registerService()
	}

	if !*flRunService {
		return false, false, nil
	}

	interactive, err := svc.IsAnInteractiveSession()
	if err != nil {
		return false, false, err
	}

	h := &handler{
		tosvc:     make(chan bool),
		fromsvc:   make(chan error),
		daemonCli: daemonCli,
	}

	var log *eventlog.Log
	if !interactive {
		log, err = eventlog.Open(*flServiceName)
		if err != nil {
			return false, false, err
		}
	}

	logrus.AddHook(&etwHook{log})
	logrus.SetOutput(io.Discard)

	service = h
	go func() {
		if interactive {
			err = debug.Run(*flServiceName, h)
		} else {
			err = svc.Run(*flServiceName, h)
		}

		h.fromsvc <- err
	}()

	// Wait for the first signal from the service handler.
	err = <-h.fromsvc
	if err != nil {
		return false, false, err
	}
	return false, true, nil
}

func (h *handler) started() error {
	// This must be delayed until daemonCli initializes Config.Root
	err := initPanicFile(filepath.Join(h.daemonCli.Config.Root, "panic.log"))
	if err != nil {
		return err
	}

	h.tosvc <- false
	return nil
}

func (h *handler) stopped(err error) {
	logrus.Debugf("Stopping service: %v", err)
	h.tosvc <- err != nil
	<-h.fromsvc
}

func (h *handler) Execute(_ []string, r <-chan svc.ChangeRequest, s chan<- svc.Status) (bool, uint32) {
	s <- svc.Status{State: svc.StartPending, Accepts: 0}
	// Unblock initService()
	h.fromsvc <- nil

	// Wait for initialization to complete.
	failed := <-h.tosvc
	if failed {
		logrus.Debug("Aborting service start due to failure during initialization")
		return true, 1
	}

	s <- svc.Status{State: svc.Running, Accepts: svc.AcceptStop | svc.AcceptShutdown | svc.Accepted(windows.SERVICE_ACCEPT_PARAMCHANGE)}
	logrus.Debug("Service running")
Loop:
	for {
		select {
		case failed = <-h.tosvc:
			break Loop
		case c := <-r:
			switch c.Cmd {
			case svc.Cmd(windows.SERVICE_CONTROL_PARAMCHANGE):
				h.daemonCli.reloadConfig()
			case svc.Interrogate:
				s <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s <- svc.Status{State: svc.StopPending, Accepts: 0}
				h.daemonCli.stop()
			}
		}
	}

	removePanicFile()
	if failed {
		return true, 1
	}
	return false, 0
}

func initPanicFile(path string) error {
	var err error
	panicFile, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o200)
	if err != nil {
		return err
	}

	st, err := panicFile.Stat()
	if err != nil {
		return err
	}

	// If there are contents in the file already, move the file out of the way
	// and replace it.
	if st.Size() > 0 {
		panicFile.Close()
		os.Rename(path, path+".old")
		panicFile, err = os.Create(path)
		if err != nil {
			return err
		}
	}

	// Update STD_ERROR_HANDLE to point to the panic file so that Go writes to
	// it when it panics. Remember the old stderr to restore it before removing
	// the panic file.
	h, err := windows.GetStdHandle(windows.STD_ERROR_HANDLE)
	if err != nil {
		return err
	}
	oldStderr = h

	err = windows.SetStdHandle(windows.STD_ERROR_HANDLE, windows.Handle(panicFile.Fd()))
	if err != nil {
		return err
	}

	// Reset os.Stderr to the panic file (so fmt.Fprintf(os.Stderr,...) actually gets redirected)
	os.Stderr = os.NewFile(panicFile.Fd(), "/dev/stderr")

	// Force threads that panic to write to stderr (the panicFile handle now), otherwise it will go into the ether
	log.SetOutput(os.Stderr)

	return nil
}

func removePanicFile() {
	if st, err := panicFile.Stat(); err == nil {
		if st.Size() == 0 {
			windows.SetStdHandle(windows.STD_ERROR_HANDLE, oldStderr)
			panicFile.Close()
			os.Remove(panicFile.Name())
		}
	}
}
<|endoftext|>