<commit_before>// Copyright 2012 The Gorilla 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 pat

import (
	"net/http"
	"net/url"
	"path"
	"strings"

	"github.com/gorilla/context"
	"github.com/gorilla/mux"
)

// New returns a new router.
func New() *Router {
	return &Router{}
}

// Router is a request router that implements a pat-like API.
//
// pat docs: http://godoc.org/github.com/bmizerany/pat
type Router struct {
	mux.Router
}

// Add registers a pattern with a handler for the given request method.
func (r *Router) Add(meth, pat string, h http.Handler) *mux.Route {
	return r.NewRoute().PathPrefix(pat).Handler(h).Methods(meth)
}

// Delete registers a pattern with a handler for DELETE requests.
func (r *Router) Delete(pat string, h http.HandlerFunc) *mux.Route {
	return r.Add("DELETE", pat, h)
}

// Get registers a pattern with a handler for GET requests.
func (r *Router) Get(pat string, h http.HandlerFunc) *mux.Route {
	return r.Add("GET", pat, h)
}

// Post registers a pattern with a handler for POST requests.
func (r *Router) Post(pat string, h http.HandlerFunc) *mux.Route {
	return r.Add("POST", pat, h)
}

// Put registers a pattern with a handler for PUT requests.
func (r *Router) Put(pat string, h http.HandlerFunc) *mux.Route {
	return r.Add("PUT", pat, h)
}

// ServeHTTP dispatches the handler registered in the matched route.
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// Clean path to canonical form and redirect.
	if p := cleanPath(req.URL.Path); p != req.URL.Path {
		w.Header().Set("Location", p)
		w.WriteHeader(http.StatusMovedPermanently)
		return
	}
	var match mux.RouteMatch
	var handler http.Handler
	if matched := r.Match(req, &match); matched {
		handler = match.Handler
		registerVars(req, match.Vars)
	}
	if handler == nil {
		if r.NotFoundHandler == nil {
			r.NotFoundHandler = http.NotFoundHandler()
		}
		handler = r.NotFoundHandler
	}
	defer context.Clear(req)
	handler.ServeHTTP(w, req)
}

// registerVars adds the matched route variables to the URL query.
func registerVars(r *http.Request, vars map[string]string) {
	parts, i := make([]string, len(vars)), 0
	for key, value := range vars {
		parts[i] = url.QueryEscape(":"+key) + "=" + url.QueryEscape(value)
		i++
	}
	q := strings.Join(parts, "&")
	if r.URL.RawQuery == "" {
		r.URL.RawQuery = q
	} else {
		r.URL.RawQuery += "&" + q
	}
}

// cleanPath returns the canonical path for p, eliminating . and .. elements.
// Borrowed from the net/http package.
func cleanPath(p string) string {
	if p == "" {
		return "/"
	}
	if p[0] != '/' {
		p = "/" + p
	}
	np := path.Clean(p)
	// path.Clean removes trailing slash except for root;
	// put the trailing slash back if necessary.
	if p[len(p)-1] == '/' && np != "/" {
		np += "/"
	}
	return np
}
<commit_msg>Added check to see if KeepContext is true and context should be kept after ending the request.<commit_after>// Copyright 2012 The Gorilla 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 pat

import (
	"net/http"
	"net/url"
	"path"
	"strings"

	"github.com/gorilla/context"
	"github.com/gorilla/mux"
)

// New returns a new router.
func New() *Router {
	return &Router{}
}

// Router is a request router that implements a pat-like API.
//
// pat docs: http://godoc.org/github.com/bmizerany/pat
type Router struct {
	mux.Router
}

// Add registers a pattern with a handler for the given request method.
func (r *Router) Add(meth, pat string, h http.Handler) *mux.Route {
	return r.NewRoute().PathPrefix(pat).Handler(h).Methods(meth)
}

// Delete registers a pattern with a handler for DELETE requests.
func (r *Router) Delete(pat string, h http.HandlerFunc) *mux.Route {
	return r.Add("DELETE", pat, h)
}

// Get registers a pattern with a handler for GET requests.
func (r *Router) Get(pat string, h http.HandlerFunc) *mux.Route {
	return r.Add("GET", pat, h)
}

// Post registers a pattern with a handler for POST requests.
func (r *Router) Post(pat string, h http.HandlerFunc) *mux.Route {
	return r.Add("POST", pat, h)
}

// Put registers a pattern with a handler for PUT requests.
func (r *Router) Put(pat string, h http.HandlerFunc) *mux.Route {
	return r.Add("PUT", pat, h)
}

// ServeHTTP dispatches the handler registered in the matched route.
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// Clean path to canonical form and redirect.
	if p := cleanPath(req.URL.Path); p != req.URL.Path {
		w.Header().Set("Location", p)
		w.WriteHeader(http.StatusMovedPermanently)
		return
	}
	var match mux.RouteMatch
	var handler http.Handler
	if matched := r.Match(req, &match); matched {
		handler = match.Handler
		registerVars(req, match.Vars)
	}
	if handler == nil {
		if r.NotFoundHandler == nil {
			r.NotFoundHandler = http.NotFoundHandler()
		}
		handler = r.NotFoundHandler
	}
	if !r.KeepContext {
		defer context.Clear(req)
	}
	handler.ServeHTTP(w, req)
}

// registerVars adds the matched route variables to the URL query.
func registerVars(r *http.Request, vars map[string]string) {
	parts, i := make([]string, len(vars)), 0
	for key, value := range vars {
		parts[i] = url.QueryEscape(":"+key) + "=" + url.QueryEscape(value)
		i++
	}
	q := strings.Join(parts, "&")
	if r.URL.RawQuery == "" {
		r.URL.RawQuery = q
	} else {
		r.URL.RawQuery += "&" + q
	}
}

// cleanPath returns the canonical path for p, eliminating . and .. elements.
// Borrowed from the net/http package.
func cleanPath(p string) string {
	if p == "" {
		return "/"
	}
	if p[0] != '/' {
		p = "/" + p
	}
	np := path.Clean(p)
	// path.Clean removes trailing slash except for root;
	// put the trailing slash back if necessary.
	if p[len(p)-1] == '/' && np != "/" {
		np += "/"
	}
	return np
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
)

const (
	ISODate          string = "2006-01-02"
	ISODateTime      string = "2006-01-02T15:04:05Z07:00"
	ISODateTimestamp string = "2006-01-02T15:04:05.000Z07:00"
	UTCDateTime      string = "UTC"
	UTCDateTimestamp string = "UTCTimestamp"
)

func isISODate(sval string) bool {
	_, err := string2date(sval, ISODate)

	if err != nil {
		return false
	}

	return true
}

func isISODateTime(sval string) bool {
	_, err := string2date(sval, ISODateTime)

	if err != nil {
		return false
	}

	return true
}

func dateFromISODateTime(sval string) (time.Time, error) {
	return string2date(sval, ISODateTime)
}

func date2string(val time.Time, format string) string {
	switch format {
	case ISODate, ISODateTime, ISODateTimestamp:
		return val.Format(format)
	case UTCDateTime:
		return val.UTC().Format(ISODateTime)
	case UTCDateTimestamp:
		return val.UTC().Format(ISODateTimestamp)
	default:
		return ""
	}
}

func string2date(sval string, format string) (time.Time, error) {
	switch format {
	case ISODate, ISODateTime, ISODateTimestamp, UTCDateTime:
		t, err := time.Parse(format, sval)
		if err != nil {
			return time.Now(), err
		}
		return t, nil
	default:
		return time.Now(), fmt.Errorf("Unknown datetime format \"%s\"", format)
	}
}
<commit_msg>Utilities to convert dates to client time<commit_after>package main

import (
	"fmt"
	"net/http"
	"time"
)

const (
	ISODate           string = "2006-01-02"
	ISODateTime       string = "2006-01-02 15:04:05"
	ISODateTimestamp  string = "2006-01-02 15:04:05.000"
	ISODateTimeZ      string = "2006-01-02 15:04:05Z07:00"
	ISODateTimestampZ string = "2006-01-02 15:04:05.000Z07:00"
	DMY               string = "02/01/2006"
	DMYTime           string = "02/01/2006 15:04:05"
	UTCDateTime       string = "UTC"
	UTCDateTimestamp  string = "UTCTimestamp"
	DateOffset        string = "Z07:00"
)

func isISODate(sval string) bool {
	_, err := string2date(sval, ISODate)

	if err != nil {
		return false
	}

	return true
}

func isISODateTime(sval string) bool {
	_, err := string2date(sval, ISODateTime)

	if err != nil {
		return false
	}

	return true
}

func dateFromISODateTime(sval string) (time.Time, error) {
	return string2date(sval, ISODateTime)
}

func date2string(val time.Time, format string) string {
	switch format {
	case ISODate, ISODateTime, ISODateTimestamp, ISODateTimeZ, ISODateTimestampZ, DMY, DMYTime:
		return val.Format(format)
	case UTCDateTime:
		return val.UTC().Format(ISODateTimeZ)
	case UTCDateTimestamp:
		return val.UTC().Format(ISODateTimestampZ)
	default:
		return ""
	}
}

func string2date(sval string, format string) (time.Time, error) {
	switch format {
	case ISODate, ISODateTime, ISODateTimestamp, ISODateTimeZ, ISODateTimestampZ, DMY, DMYTime, DateOffset:
		t, err := time.Parse(format, sval)
		if err != nil {
			return time.Now(), err
		}
		return t, nil
	default:
		return time.Now(), fmt.Errorf("Unknown datetime format \"%s\"", format)
	}
}

func server2ClientDmy(r *http.Request, serverTime time.Time) string {
	t := server2ClientLocal(r, serverTime)
	return date2string(t, DMY)
}

func server2ClientDmyTime(r *http.Request, serverTime time.Time) string {
	t := server2ClientLocal(r, serverTime)
	return date2string(t, DMYTime)
}

func server2ClientLocal(r *http.Request, serverTime time.Time) time.Time {
	timeOffset := 0

	cookie, err := r.Cookie("time_zone_offset")
	if err != nil && err != http.ErrNoCookie {
		return serverTime.UTC()
	} else if err == http.ErrNoCookie {
		timeOffset = 0
	} else {
		timeOffset = string2int(cookie.Value)
	}

	return serverTime.UTC().Add(time.Duration(-1*timeOffset) * time.Minute)
}
<|endoftext|>


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

// Package eventlog provides an interface to the Windows Event Log.
package eventlog

import (
	"errors"
	"fmt"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/sys/windows"
	"github.com/google/winops/winlog/wevtapi"
)

// Handle maps a handle to an event log resource (EVT_HANDLE). Close() must be called to release the handle.
//
// Note that the order in which handles are closed may matter. Parent handles should not be closed until all
// uses of the handles (queries, etc) are complete.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtclose
type Handle struct {
	handle windows.Handle
}

// Close releases a Handle.
func (h *Handle) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// An Event is a Handle to an event.
type Event Handle

// Close releases an Event.
func (h *Event) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// A RenderContext is a Handle which tracks a Context as returned by EvtCreateRenderContext.
type RenderContext Handle

// Close releases a RenderContext.
func (h *RenderContext) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// A ResultSet is a Handle returned by a Query or Subscription
type ResultSet Handle

// Close releases a ResultSet.
func (h *ResultSet) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// A Session is a Handle returned by OpenSession
type Session Handle

// Close releases a Session.
func (h *Session) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// EvtRenderContextFlags specify which types of values to render from a given event.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ne-winevt-evt_render_context_flags
type EvtRenderContextFlags uint32

const (
	// EvtRenderContextValues renders specific properties from the event.
	EvtRenderContextValues EvtRenderContextFlags = iota
	// EvtRenderContextSystem renders the system properties under the System element.
	EvtRenderContextSystem
	// EvtRenderContextUser renders all user-defined properties under the UserData or EventData element.
	EvtRenderContextUser
)

// CreateRenderContext creates a context that specifies the information in the event that you want to render.
//
// The RenderContext is used to obtain only a subset of event data when querying events.
// Without a RenderContext, the entirety of the log data will be returned.
//
// Passing one of EvtRenderContextSystem or EvtRenderContextUser (with valuePaths nil)
// will render all properties under the corresponding element (System or User). Passing
// EvtRenderContextValues along with a list of valuePaths allows the caller to obtain individual
// event elements. valuePaths must be well formed XPath expressions. See the documentation
// for EvtCreateRenderContext and EVT_RENDER_CONTEXT_FLAGS for more detail.
//
// Example, rendering all System values:
//		eventlog.CreateRenderContext(eventlog.EvtRenderContextSystem, nil)
//
// Example, rendering specific values:
//		eventlog.CreateRenderContext(eventlog.EvtRenderContextValues, &[]string{
//				"Event/System/TimeCreated/@SystemTime", "Event/System/Provider/@Name"})
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtcreaterendercontext
func CreateRenderContext(flags EvtRenderContextFlags, valuePaths *[]string) (RenderContext, error) {
	rc := RenderContext{}

	pathsPtr := uintptr(0)
	p := []*uint16{}
	if valuePaths != nil {
		for _, v := range *valuePaths {
			ptr, err := syscall.UTF16PtrFromString(v)
			if err != nil {
				return rc, fmt.Errorf("syscall.UTF16PtrFromString(%s): %w", v, err)
			}
			p = append(p, ptr)
		}
		pathsPtr = uintptr(unsafe.Pointer(&p[0]))
	}

	var err error
	rc.handle, err = wevtapi.EvtCreateRenderContext(uint32(len(p)), uintptr(pathsPtr), uint32(flags))
	return rc, err
}

// An EventSet holds one or more event handles.
//
// Close() must be called to release the event handles when finished.
type EventSet struct {
	Events []Event
	Count  uint32
}

// Close releases all events in the EventSet.
func (e *EventSet) Close() {
	for _, evt := range e.Events {
		evt.Close()
	}
}

// Next gets the next event(s) returned by a query or subscription.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtnext
func Next(handle ResultSet, count uint32, timeout *time.Duration) (EventSet, error) {
	es := EventSet{}

	defaultTimeout := 2000 * time.Millisecond
	if timeout == nil {
		timeout = &defaultTimeout
	}

	// Get handles to events from the result set.
	evts := make([]windows.Handle, count)
	err := wevtapi.EvtNext(
		handle.handle,                  // Handle to query or subscription result set.
		count,                          // The number of events to attempt to retrieve.
		&evts[0],                       // Pointer to the array of event handles.
		uint32(timeout.Milliseconds()), // Timeout in milliseconds to wait.
		0,                              // Reserved. Must be zero.
		&es.Count)                      // The number of handles in the array that are set by the API.
	if err == windows.ERROR_NO_MORE_ITEMS {
		return es, err
	} else if err != nil {
		return es, fmt.Errorf("wevtapi.EvtNext: %w", err)
	}

	for i := 0; i < int(es.Count); i++ {
		es.Events = append(es.Events, Event{handle: evts[i]})
	}

	return es, nil
}

// Query runs a query to retrieve events from a channel or log file that match the specified query criteria.
//
// Session is only required for remote connections; leave as nil for the local log. Flags can be any of
// wevtapi.EVT_QUERY_FLAGS.
//
// The session handle must remain open until all subsequent processing on the query results have completed. Call
// Close() once complete.
//
// Example:
// 	 conn, err := eventlog.Query(nil, "Windows Powershell", "*", wevtapi.EvtQueryReverseDirection)
// 	 if err != nil {
//     return err
//	 }
//	 defer conn.Close()
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtquery
func Query(session *Session, path string, query string, flags uint32) (ResultSet, error) {
	var rs ResultSet
	var err error

	var s windows.Handle
	if session != nil {
		s = session.handle
	}
	rs.handle, err = wevtapi.EvtQuery(s, windows.StringToUTF16Ptr(path), windows.StringToUTF16Ptr(query), flags)
	if err != nil {
		return rs, fmt.Errorf("EvtQuery: %w", err)
	}
	if rs.handle == windows.InvalidHandle {
		return rs, errors.New("invalid query")
	}
	return rs, nil
}

// EvtVariantData models the union inside of the EVT_VARIANT structure.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ns-winevt-evt_variant
type EvtVariantData struct {
	BooleanVal    bool
	SByteVal      int8
	Int16Val      int16
	Int32Val      int32
	Int64Val      int64
	ByteVal       uint8
	UInt16Val     uint16
	UInt32Val     uint32
	UInt64Val     uint64
	SingleVal     float32
	DoubleVal     float64
	FileTimeVal   windows.Filetime
	SysTimeVal    windows.Systemtime
	GuidVal       windows.GUID
	StringVal     string
	AnsiStringVal string
	BinaryVal     byte
	SidVal        windows.SID
	SizeTVal      uint32
	BooleanArr    *[]bool
	SByteArr      *[]int8
	Int16Arr      *[]int16
	Int32Arr      *[]int32
	Int64Arr      *[]int64
	ByteArr       *[]uint16
	UInt16Arr     *[]uint16
	UInt32Arr     *[]uint32
	UInt64Arr     *[]uint64
	SingleArr     *[]float32
	DoubleArr     *[]float64
	FileTimeArr   *[]windows.Filetime
	SysTimeArr    *[]windows.Systemtime
	GuidArr       *[]windows.GUID
	StringArr     *[]string
	AnsiStringArr *[]string
	SidArr        *[]windows.SID
	SizeTArr      *[]uint32
	EvtHandleVal  windows.Handle
	XmlVal        string
	XmlValArr     *[]string
}

// EvtVariantType(EVT_VARIANT_TYPE) defines the possible data types of a EVT_VARIANT data item.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ne-winevt-evt_variant_type
type EvtVariantType uint32

const (
	EvtVarTypeNull EvtVariantType = iota
	EvtVarTypeString
	EvtVarTypeAnsiString
	EvtVarTypeSByte
	EvtVarTypeByte
	EvtVarTypeInt16
	EvtVarTypeUInt16
	EvtVarTypeInt32
	EvtVarTypeUInt32
	EvtVarTypeInt64
	EvtVarTypeUInt64
	EvtVarTypeSingle
	EvtVarTypeDouble
	EvtVarTypeBoolean
	EvtVarTypeBinary
	EvtVarTypeGuid
	EvtVarTypeSizeT
	EvtVarTypeFileTime
	EvtVarTypeSysTime
	EvtVarTypeSid
	EvtVarTypeHexInt32
	EvtVarTypeHexInt64
	EvtVarTypeEvtHandle
	EvtVarTypeEvtXml
)

// EvtVariant (EVT_VARIANT) contains event data or property values.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ns-winevt-evt_variant
type EvtVariant struct {
	Count uint32
	Type  EvtVariantType
	Data  EvtVariantData
}

// Fragment describes a renderable fragment; an event or to a bookmark.
type Fragment interface {
	Handle() windows.Handle
}

// Render renders a fragment (bookmark or event) as an XML string.
//
// This function renders the entire fragment as XML. To render only specific elements of the event, use RenderValues.
//
// Flags can be either EvtRenderEventValues or EvtRenderEventXml.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtrender
func Render(fragment Fragment, flag uint32) (string, error) {
	var bufferUsed uint32
	var propertyCount uint32

	if flag == wevtapi.EvtRenderEventValues {
		return "", fmt.Errorf("EvtRenderEventValues requires the RenderValues function")
	}

	// Call EvtRender with a null buffer to get the required buffer size.
	err := wevtapi.EvtRender(
		0,
		fragment.Handle(),
		flag,
		0,
		nil,
		&bufferUsed,
		&propertyCount)
	if err != syscall.ERROR_INSUFFICIENT_BUFFER {
		return "", fmt.Errorf("wevtapi.EvtRender: %w", err)
	}

	// Create a buffer based on the buffer size required.
	buf := make([]uint16, bufferUsed/2)

	// Render the fragment according to the flag.
	if err = wevtapi.EvtRender(
		0,
		fragment.Handle(),
		flag,
		bufferUsed,
		unsafe.Pointer(&buf[0]),
		&bufferUsed,
		&propertyCount); err != nil {
		return "", fmt.Errorf("wevtapi.EvtRender: %w", err)
	}

	return syscall.UTF16ToString(buf), nil
}
<commit_msg>Add RenderValues<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package eventlog provides an interface to the Windows Event Log.
package eventlog

import (
	"errors"
	"fmt"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/sys/windows"
	"github.com/google/logger"
	"github.com/google/winops/winlog/wevtapi"
)

// Handle maps a handle to an event log resource (EVT_HANDLE). Close() must be called to release the handle.
//
// Note that the order in which handles are closed may matter. Parent handles should not be closed until all
// uses of the handles (queries, etc) are complete.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtclose
type Handle struct {
	handle windows.Handle
}

// Close releases a Handle.
func (h *Handle) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// An Event is a Handle to an event.
type Event Handle

// Close releases an Event.
func (h *Event) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// A RenderContext is a Handle which tracks a Context as returned by EvtCreateRenderContext.
type RenderContext Handle

// Close releases a RenderContext.
func (h *RenderContext) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// A ResultSet is a Handle returned by a Query or Subscription
type ResultSet Handle

// Close releases a ResultSet.
func (h *ResultSet) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// A Session is a Handle returned by OpenSession
type Session Handle

// Close releases a Session.
func (h *Session) Close() {
	if h != nil {
		wevtapi.EvtClose(h.handle)
	}
}

// EvtRenderContextFlags specify which types of values to render from a given event.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ne-winevt-evt_render_context_flags
type EvtRenderContextFlags uint32

const (
	// EvtRenderContextValues renders specific properties from the event.
	EvtRenderContextValues EvtRenderContextFlags = iota
	// EvtRenderContextSystem renders the system properties under the System element.
	EvtRenderContextSystem
	// EvtRenderContextUser renders all user-defined properties under the UserData or EventData element.
	EvtRenderContextUser
)

// CreateRenderContext creates a context that specifies the information in the event that you want to render.
//
// The RenderContext is used to obtain only a subset of event data when querying events.
// Without a RenderContext, the entirety of the log data will be returned.
//
// Passing one of EvtRenderContextSystem or EvtRenderContextUser (with valuePaths nil)
// will render all properties under the corresponding element (System or User). Passing
// EvtRenderContextValues along with a list of valuePaths allows the caller to obtain individual
// event elements. valuePaths must be well formed XPath expressions. See the documentation
// for EvtCreateRenderContext and EVT_RENDER_CONTEXT_FLAGS for more detail.
//
// Example, rendering all System values:
//		eventlog.CreateRenderContext(eventlog.EvtRenderContextSystem, nil)
//
// Example, rendering specific values:
//		eventlog.CreateRenderContext(eventlog.EvtRenderContextValues, &[]string{
//				"Event/System/TimeCreated/@SystemTime", "Event/System/Provider/@Name"})
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtcreaterendercontext
func CreateRenderContext(flags EvtRenderContextFlags, valuePaths *[]string) (RenderContext, error) {
	rc := RenderContext{}

	pathsPtr := uintptr(0)
	p := []*uint16{}
	if valuePaths != nil {
		for _, v := range *valuePaths {
			ptr, err := syscall.UTF16PtrFromString(v)
			if err != nil {
				return rc, fmt.Errorf("syscall.UTF16PtrFromString(%s): %w", v, err)
			}
			p = append(p, ptr)
		}
		pathsPtr = uintptr(unsafe.Pointer(&p[0]))
	}

	var err error
	rc.handle, err = wevtapi.EvtCreateRenderContext(uint32(len(p)), uintptr(pathsPtr), uint32(flags))
	return rc, err
}

// An EventSet holds one or more event handles.
//
// Close() must be called to release the event handles when finished.
type EventSet struct {
	Events []Event
	Count  uint32
}

// Close releases all events in the EventSet.
func (e *EventSet) Close() {
	for _, evt := range e.Events {
		evt.Close()
	}
}

// Next gets the next event(s) returned by a query or subscription.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtnext
func Next(handle ResultSet, count uint32, timeout *time.Duration) (EventSet, error) {
	es := EventSet{}

	defaultTimeout := 2000 * time.Millisecond
	if timeout == nil {
		timeout = &defaultTimeout
	}

	// Get handles to events from the result set.
	evts := make([]windows.Handle, count)
	err := wevtapi.EvtNext(
		handle.handle,                  // Handle to query or subscription result set.
		count,                          // The number of events to attempt to retrieve.
		&evts[0],                       // Pointer to the array of event handles.
		uint32(timeout.Milliseconds()), // Timeout in milliseconds to wait.
		0,                              // Reserved. Must be zero.
		&es.Count)                      // The number of handles in the array that are set by the API.
	if err == windows.ERROR_NO_MORE_ITEMS {
		return es, err
	} else if err != nil {
		return es, fmt.Errorf("wevtapi.EvtNext: %w", err)
	}

	for i := 0; i < int(es.Count); i++ {
		es.Events = append(es.Events, Event{handle: evts[i]})
	}

	return es, nil
}

// Query runs a query to retrieve events from a channel or log file that match the specified query criteria.
//
// Session is only required for remote connections; leave as nil for the local log. Flags can be any of
// wevtapi.EVT_QUERY_FLAGS.
//
// The session handle must remain open until all subsequent processing on the query results have completed. Call
// Close() once complete.
//
// Example:
// 	 conn, err := eventlog.Query(nil, "Windows Powershell", "*", wevtapi.EvtQueryReverseDirection)
// 	 if err != nil {
//     return err
//	 }
//	 defer conn.Close()
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtquery
func Query(session *Session, path string, query string, flags uint32) (ResultSet, error) {
	var rs ResultSet
	var err error

	var s windows.Handle
	if session != nil {
		s = session.handle
	}
	rs.handle, err = wevtapi.EvtQuery(s, windows.StringToUTF16Ptr(path), windows.StringToUTF16Ptr(query), flags)
	if err != nil {
		return rs, fmt.Errorf("EvtQuery: %w", err)
	}
	if rs.handle == windows.InvalidHandle {
		return rs, errors.New("invalid query")
	}
	return rs, nil
}

// EvtVariantData models the union inside of the EVT_VARIANT structure.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ns-winevt-evt_variant
type EvtVariantData struct {
	BooleanVal    bool
	SByteVal      int8
	Int16Val      int16
	Int32Val      int32
	Int64Val      int64
	ByteVal       uint8
	UInt16Val     uint16
	UInt32Val     uint32
	UInt64Val     uint64
	SingleVal     float32
	DoubleVal     float64
	FileTimeVal   windows.Filetime
	SysTimeVal    windows.Systemtime
	GuidVal       windows.GUID
	StringVal     string
	AnsiStringVal string
	BinaryVal     byte
	SidVal        windows.SID
	SizeTVal      uint32
	BooleanArr    *[]bool
	SByteArr      *[]int8
	Int16Arr      *[]int16
	Int32Arr      *[]int32
	Int64Arr      *[]int64
	ByteArr       *[]uint16
	UInt16Arr     *[]uint16
	UInt32Arr     *[]uint32
	UInt64Arr     *[]uint64
	SingleArr     *[]float32
	DoubleArr     *[]float64
	FileTimeArr   *[]windows.Filetime
	SysTimeArr    *[]windows.Systemtime
	GuidArr       *[]windows.GUID
	StringArr     *[]string
	AnsiStringArr *[]string
	SidArr        *[]windows.SID
	SizeTArr      *[]uint32
	EvtHandleVal  windows.Handle
	XmlVal        string
	XmlValArr     *[]string
}

// EvtVariantType(EVT_VARIANT_TYPE) defines the possible data types of a EVT_VARIANT data item.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ne-winevt-evt_variant_type
type EvtVariantType uint32

const (
	EvtVarTypeNull EvtVariantType = iota
	EvtVarTypeString
	EvtVarTypeAnsiString
	EvtVarTypeSByte
	EvtVarTypeByte
	EvtVarTypeInt16
	EvtVarTypeUInt16
	EvtVarTypeInt32
	EvtVarTypeUInt32
	EvtVarTypeInt64
	EvtVarTypeUInt64
	EvtVarTypeSingle
	EvtVarTypeDouble
	EvtVarTypeBoolean
	EvtVarTypeBinary
	EvtVarTypeGuid
	EvtVarTypeSizeT
	EvtVarTypeFileTime
	EvtVarTypeSysTime
	EvtVarTypeSid
	EvtVarTypeHexInt32
	EvtVarTypeHexInt64
	EvtVarTypeEvtHandle
	EvtVarTypeEvtXml
)

// EvtVariant (EVT_VARIANT) contains event data or property values.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ns-winevt-evt_variant
type EvtVariant struct {
	Count uint32
	Type  EvtVariantType
	Data  EvtVariantData
}

// Fragment describes a renderable fragment; an event or to a bookmark.
type Fragment interface {
	Handle() windows.Handle
}

// Render renders a fragment (bookmark or event) as an XML string.
//
// This function renders the entire fragment as XML. To render only specific elements of the event, use RenderValues.
//
// Flags can be either EvtRenderEventValues or EvtRenderEventXml.
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtrender
func Render(fragment Fragment, flag uint32) (string, error) {
	var bufferUsed uint32
	var propertyCount uint32

	if flag == wevtapi.EvtRenderEventValues {
		return "", fmt.Errorf("EvtRenderEventValues requires the RenderValues function")
	}

	// Call EvtRender with a null buffer to get the required buffer size.
	err := wevtapi.EvtRender(
		0,
		fragment.Handle(),
		flag,
		0,
		nil,
		&bufferUsed,
		&propertyCount)
	if err != syscall.ERROR_INSUFFICIENT_BUFFER {
		return "", fmt.Errorf("wevtapi.EvtRender: %w", err)
	}

	// Create a buffer based on the buffer size required.
	buf := make([]uint16, bufferUsed/2)

	// Render the fragment according to the flag.
	if err = wevtapi.EvtRender(
		0,
		fragment.Handle(),
		flag,
		bufferUsed,
		unsafe.Pointer(&buf[0]),
		&bufferUsed,
		&propertyCount); err != nil {
		return "", fmt.Errorf("wevtapi.EvtRender: %w", err)
	}

	return syscall.UTF16ToString(buf), nil
}

// RenderValues renders specific elements from a fragment (event).
//
// You must supply a RenderContext from CreateRenderContext. The RenderContext determines which values are rendered from the fragment.
//
// The rendered events are returned by Windows as variants (EVT_VARIANT) in a wide variety of types. We do our best to
// cast these into Go types, and return the results encapsulated in an EvtVariantData. EvtVariantData holds all possible types, but
// only the rendered value should be non-nil on return. The EvtVariant.Type field will indicate which of the EvtVariantData fields
// should hold the rendered data.
//
// For example, rendering a string fragment successfully should return:
//		EvtVariant {
//			Type: EvtVarTypeString
//			Data: EvtVariantData {
//				...
//				StringVal: "my rendered string"
//				...
//			}
//		}
//
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/nf-winevt-evtrender
// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ns-winevt-evt_variant
func RenderValues(renderCtx RenderContext, fragment Fragment) ([]EvtVariant, error) {
	var bufferUsed uint32
	var propertyCount uint32
	var vals []EvtVariant

	// Call EvtRender with a null buffer to get the required buffer size.
	err := wevtapi.EvtRender(
		renderCtx.handle,
		fragment.Handle(),
		wevtapi.EvtRenderEventValues,
		0,
		nil,
		&bufferUsed,
		&propertyCount)
	if err != syscall.ERROR_INSUFFICIENT_BUFFER {
		return nil, fmt.Errorf("wevtapi.EvtRender: %w", err)
	}

	// Create a buffer to hold the EVT_VARIANT objects returned by the query.
	//
	// Ref: https://docs.microsoft.com/en-us/windows/win32/api/winevt/ns-winevt-evt_variant
	buf := make([]struct {
		Raw   [8]byte // Go represents the union as a byte slice, sized based on the largest element in the union.
		Count uint32
		Type  uint32
	}, propertyCount)

	// Render the fragment according to the flag.
	if err = wevtapi.EvtRender(
		renderCtx.handle,
		fragment.Handle(),
		wevtapi.EvtRenderEventValues,
		bufferUsed,
		unsafe.Pointer(&buf[0]),
		&bufferUsed,
		&propertyCount); err != nil {
		return nil, fmt.Errorf("wevtapi.EvtRender: %w", err)
	}

	// The EVT_VARIANT union can be holding any of the union's supported data types.
	// To make it useable, we look for the type in the Type field and cast accordingly.
	for i := 0; i < int(propertyCount); i++ {
		v := EvtVariant{Type: EvtVariantType(buf[i].Type)}
		switch buf[i].Type {
		case uint32(EvtVarTypeNull):
			continue
		case uint32(EvtVarTypeString):
			ptr := *(**uint16)(unsafe.Pointer(&buf[i].Raw))
			v.Data.StringVal = windows.UTF16PtrToString(ptr)
		case uint32(EvtVarTypeSByte):
			v.Data.SByteVal = *(*int8)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeByte):
			v.Data.ByteVal = *(*uint8)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeInt16):
			v.Data.Int16Val = *(*int16)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeInt32), uint32(EvtVarTypeHexInt32):
			v.Data.Int32Val = *(*int32)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeInt64), uint32(EvtVarTypeHexInt64):
			v.Data.Int64Val = *(*int64)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeUInt16):
			v.Data.UInt16Val = *(*uint16)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeUInt32):
			v.Data.UInt32Val = *(*uint32)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeUInt64):
			v.Data.UInt64Val = *(*uint64)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeBoolean):
			v.Data.BooleanVal = *(*bool)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeGuid):
			v.Data.GuidVal = *(*windows.GUID)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeFileTime):
			v.Data.FileTimeVal = *(*windows.Filetime)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeSid):
			v.Data.SidVal = *(*windows.SID)(unsafe.Pointer(&buf[i].Raw))
		case uint32(EvtVarTypeSysTime):
			v.Data.SysTimeVal = *(*windows.Systemtime)(unsafe.Pointer(&buf[i].Raw))
		default:
			logger.Warningf("Unsupported type: %v", buf[i].Type)
		}
		vals = append(vals, v)
	}

	return vals, nil
}
<|endoftext|>


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

package storage

import (
	"errors"
	"fmt"
	"strconv"

	"github.com/google/logger"
	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
)

// Partition represents a MSFT_Partition object.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition
type Partition struct {
	DiskNumber           int32
	PartitionNumber      int32
	DriveLetter          string
	AccessPaths          []string
	OperationalStatus    int32
	TransitionState      int32
	Size                 uint64
	MbrType              int32
	GptType              string
	GUID                 string
	IsReadOnly           bool
	IsOffline            bool
	IsSystem             bool
	IsBoot               bool
	IsActive             bool
	IsHidden             bool
	IsShadowCopy         bool
	NoDefaultDriveLetter bool

	handle *ole.IDispatch
}

// PartitionSupportedSize represents the minimum and maximum sizes a partition can be resized to.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-getsupportedsizes
type PartitionSupportedSize struct {
	SizeMin uint64
	SizeMax uint64
}

// Close releases the handle to the partition.
func (p *Partition) Close() {
	if p.handle != nil {
		p.handle.Release()
	}
}

// Delete attempts to delete a partition.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-deleteobject
func (p *Partition) Delete() (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	resultRaw, err := oleutil.CallMethod(p.handle, "DeleteObject", &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("DeleteObject: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during deletion: %d", val)
	}
	return stat, nil
}

// Offline takes the partition offline.
//
// Example:
//		p.Offline()
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-offline
func (p *Partition) Offline() (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	res, err := oleutil.CallMethod(p.handle, "Offline", &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("Offline(): %w", err)
	} else if val, ok := res.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during offline: %d", val)
	}
	return stat, nil
}

// Online brings the partition online by mounting the associated volume (if one exists).
//
// Example:
//		p.Online()
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-online
func (p *Partition) Online() (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	res, err := oleutil.CallMethod(p.handle, "Online", &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("Online(): %w", err)
	} else if val, ok := res.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during online: %d", val)
	}
	return stat, nil
}

// AddAccessPath adds a mount path or drive letter assignment to the partition.
//
// Example: assign a Drive letter with D:
//		p.AddAccessPath("D:", false)
//
// Example: Automatically assign the next available Drive Letter:
//		p.AddAccessPath("", true)
//
// Note: You cannot specify both a valid drive letter and auto assignment as true together.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/addaccesspath-msft-partition
func (p *Partition) AddAccessPath(accessPath string, autoAssign bool) (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	var resultRaw *ole.VARIANT
	var err error
	if autoAssign {
		resultRaw, err = oleutil.CallMethod(p.handle, "AddAccessPath", nil, autoAssign, &extendedStatus)
	} else {
		resultRaw, err = oleutil.CallMethod(p.handle, "AddAccessPath", accessPath, nil, &extendedStatus)
	}
	if err != nil {
		return stat, fmt.Errorf("AddAccessPath: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during AddAccessPath: %d", val)
	}
	return stat, nil
}

// RemoveAccessPath removes the access path from the partition.
//
// Example: Remove the driveLetter of D: from a partition
//		p.RemoveAccessPath("D:")
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/removeaccesspath-msft-partition
func (p *Partition) RemoveAccessPath(accessPath string) (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	resultRaw, err := oleutil.CallMethod(p.handle, "RemoveAccessPath", accessPath, &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("RemoveAccessPath: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during RemoveAccessPath: %d", val)
	}
	return stat, nil
}

// Query reads and populates the partition state.
func (p *Partition) Query() error {
	if p.handle == nil {
		return fmt.Errorf("invalid handle")
	}

	// DriveLetter
	prop, err := oleutil.GetProperty(p.handle, "DriveLetter")
	if err != nil {
		return fmt.Errorf("oleutil.GetProperty(DriveLetter): %w", err)
	}
	// DriveLetter is represented as Char16 (Ascii)
	p.DriveLetter = string(rune(prop.Val))

	// AccessPaths
	prop, err = oleutil.GetProperty(p.handle, "AccessPaths")
	if err != nil {
		return fmt.Errorf("oleutil.GetProperty(AccessPaths): %w", err)
	}
	for _, pa := range prop.ToArray().ToValueArray() {
		conv, ok := pa.(string)
		if !ok {
			return errors.New("error converting access path")
		}
		p.AccessPaths = append(p.AccessPaths, conv)
	}

	// GptType
	prop, err = oleutil.GetProperty(p.handle, "GptType")
	if err != nil {
		return fmt.Errorf("oleutil.GetProperty(GptType): %w", err)
	}
	p.GptType = prop.ToString()

	// GUID
	prop, err = oleutil.GetProperty(p.handle, "Guid")
	if err != nil {
		return fmt.Errorf("oleutil.GetProperty(Guid): %w", err)
	}
	p.GUID = prop.ToString()

	// All the non-strings
	for _, prop := range [][]interface{}{
		[]interface{}{"DiskNumber", &p.DiskNumber},
		[]interface{}{"PartitionNumber", &p.PartitionNumber},
		[]interface{}{"OperationalStatus", &p.OperationalStatus},
		[]interface{}{"TransitionState", &p.TransitionState},
		[]interface{}{"Size", &p.Size},
		[]interface{}{"MbrType", &p.MbrType},
		[]interface{}{"IsReadOnly", &p.IsReadOnly},
		[]interface{}{"IsOffline", &p.IsOffline},
		[]interface{}{"IsSystem", &p.IsSystem},
		[]interface{}{"IsBoot", &p.IsBoot},
		[]interface{}{"IsActive", &p.IsActive},
		[]interface{}{"IsHidden", &p.IsHidden},
		[]interface{}{"IsShadowCopy", &p.IsShadowCopy},
		[]interface{}{"NoDefaultDriveLetter", &p.NoDefaultDriveLetter},
	} {
		val, err := oleutil.GetProperty(p.handle, prop[0].(string))
		if err != nil {
			return fmt.Errorf("oleutil.GetProperty(%s): %w", prop[0].(string), err)
		}
		if err := assignVariant(val.Value(), prop[1]); err != nil {
			logger.Warningf("assignVariant(%s): %v", prop[0].(string), err)
		}
	}
	return nil
}

// Resize attempts to resize a partition.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-resize
func (p *Partition) Resize(size uint64) (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	// Convert the unint to a string because of because of https://docs.microsoft.com/en-us/previous-versions//aa393262%28v=vs.85%29?redirectedfrom=MSDN
	resultRaw, err := oleutil.CallMethod(p.handle, "Resize", strconv.FormatUint(size, 10), &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("Resize: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during resize: %d", val)
	}
	return stat, nil
}

// A PartitionSet contains one or more Partitions.
type PartitionSet struct {
	Partitions []Partition
}

// Close releases all Partition handles inside a PartitionSet.
func (s *PartitionSet) Close() {
	for _, p := range s.Partitions {
		p.Close()
	}
}

// GetSupportedSize retrieves the minimum and maximum sizes that the partition can be resized to
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-getsupportedsizes
func (p *Partition) GetSupportedSize() (PartitionSupportedSize, ExtendedStatus, error) {
	size := PartitionSupportedSize{}
	stat := ExtendedStatus{}

	var sizemin ole.VARIANT
	ole.VariantInit(&sizemin)
	var sizemax ole.VARIANT
	ole.VariantInit(&sizemax)
	var extendedstatus ole.VARIANT
	ole.VariantInit(&extendedstatus)

	resultRaw, err := oleutil.CallMethod(p.handle, "GetSupportedSize", &sizemin, &sizemax, &extendedstatus)
	if err != nil {
		return size, stat, fmt.Errorf("GetSupportedSize: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return size, stat, fmt.Errorf("error code returned during GetSupportedSize: %d", val)
	}

	// Convert the results from an interface to uint64
	// Results are returned as strings because of https://docs.microsoft.com/en-us/previous-versions//aa393262%28v=vs.85%29?redirectedfrom=MSDN
	size.SizeMin, err = strconv.ParseUint(sizemin.Value().(string), 10, 64)
	if err != nil {
		return size, stat, fmt.Errorf("error attempting to parse sizemin: %w", err)
	}
	size.SizeMax, err = strconv.ParseUint(sizemax.Value().(string), 10, 64)
	if err != nil {
		return size, stat, fmt.Errorf("error attempting to parse sizemax: %w", err)
	}
	return size, stat, nil
}

// GetPartitions queries for local partitions.
//
// Close() must be called on the resulting PartitionSet to ensure all disks are released.
//
// Get all partitions:
//		svc.GetPartitions("")
//
// To get specific partitions, provide a valid WMI query filter string, for example:
//		svc.GetPartitions("WHERE DiskNumber=1")
func (svc *Service) GetPartitions(filter string) (PartitionSet, error) {
	parts := PartitionSet{}
	query := "SELECT * FROM MSFT_Partition"
	if filter != "" {
		query = fmt.Sprintf("%s %s", query, filter)
	}
	raw, err := oleutil.CallMethod(svc.wmiSvc, "ExecQuery", query)
	if err != nil {
		return parts, fmt.Errorf("ExecQuery(%s): %w", query, err)
	}
	result := raw.ToIDispatch()
	defer result.Release()

	countVar, err := oleutil.GetProperty(result, "Count")
	if err != nil {
		return parts, fmt.Errorf("oleutil.GetProperty(Count): %w", err)
	}
	count := int(countVar.Val)

	for i := 0; i < count; i++ {
		part := Partition{}
		itemRaw, err := oleutil.CallMethod(result, "ItemIndex", i)
		if err != nil {
			return parts, fmt.Errorf("oleutil.CallMethod(ItemIndex, %d): %w", i, err)
		}
		part.handle = itemRaw.ToIDispatch()

		if err := part.Query(); err != nil {
			return parts, err
		}

		parts.Partitions = append(parts.Partitions, part)
	}

	return parts, nil
}
<commit_msg>Leave NUL as empty string in drive letters.<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License")
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package storage

import (
	"errors"
	"fmt"
	"strconv"

	"github.com/google/logger"
	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
)

// Partition represents a MSFT_Partition object.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition
type Partition struct {
	DiskNumber           int32
	PartitionNumber      int32
	DriveLetter          string
	AccessPaths          []string
	OperationalStatus    int32
	TransitionState      int32
	Size                 uint64
	MbrType              int32
	GptType              string
	GUID                 string
	IsReadOnly           bool
	IsOffline            bool
	IsSystem             bool
	IsBoot               bool
	IsActive             bool
	IsHidden             bool
	IsShadowCopy         bool
	NoDefaultDriveLetter bool

	handle *ole.IDispatch
}

// PartitionSupportedSize represents the minimum and maximum sizes a partition can be resized to.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-getsupportedsizes
type PartitionSupportedSize struct {
	SizeMin uint64
	SizeMax uint64
}

// Close releases the handle to the partition.
func (p *Partition) Close() {
	if p.handle != nil {
		p.handle.Release()
	}
}

// Delete attempts to delete a partition.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-deleteobject
func (p *Partition) Delete() (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	resultRaw, err := oleutil.CallMethod(p.handle, "DeleteObject", &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("DeleteObject: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during deletion: %d", val)
	}
	return stat, nil
}

// Offline takes the partition offline.
//
// Example:
//		p.Offline()
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-offline
func (p *Partition) Offline() (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	res, err := oleutil.CallMethod(p.handle, "Offline", &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("Offline(): %w", err)
	} else if val, ok := res.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during offline: %d", val)
	}
	return stat, nil
}

// Online brings the partition online by mounting the associated volume (if one exists).
//
// Example:
//		p.Online()
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-online
func (p *Partition) Online() (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	res, err := oleutil.CallMethod(p.handle, "Online", &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("Online(): %w", err)
	} else if val, ok := res.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during online: %d", val)
	}
	return stat, nil
}

// AddAccessPath adds a mount path or drive letter assignment to the partition.
//
// Example: assign a Drive letter with D:
//		p.AddAccessPath("D:", false)
//
// Example: Automatically assign the next available Drive Letter:
//		p.AddAccessPath("", true)
//
// Note: You cannot specify both a valid drive letter and auto assignment as true together.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/addaccesspath-msft-partition
func (p *Partition) AddAccessPath(accessPath string, autoAssign bool) (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	var resultRaw *ole.VARIANT
	var err error
	if autoAssign {
		resultRaw, err = oleutil.CallMethod(p.handle, "AddAccessPath", nil, autoAssign, &extendedStatus)
	} else {
		resultRaw, err = oleutil.CallMethod(p.handle, "AddAccessPath", accessPath, nil, &extendedStatus)
	}
	if err != nil {
		return stat, fmt.Errorf("AddAccessPath: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during AddAccessPath: %d", val)
	}
	return stat, nil
}

// RemoveAccessPath removes the access path from the partition.
//
// Example: Remove the driveLetter of D: from a partition
//		p.RemoveAccessPath("D:")
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/removeaccesspath-msft-partition
func (p *Partition) RemoveAccessPath(accessPath string) (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	resultRaw, err := oleutil.CallMethod(p.handle, "RemoveAccessPath", accessPath, &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("RemoveAccessPath: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during RemoveAccessPath: %d", val)
	}
	return stat, nil
}

// Query reads and populates the partition state.
func (p *Partition) Query() error {
	if p.handle == nil {
		return fmt.Errorf("invalid handle")
	}

	// DriveLetter
	prop, err := oleutil.GetProperty(p.handle, "DriveLetter")
	if err != nil {
		return fmt.Errorf("oleutil.GetProperty(DriveLetter): %w", err)
	}
	// DriveLetter is represented as Char16 (Ascii)
	if prop.Val != 0 { // leave NUL as empty string
		p.DriveLetter = string(rune(prop.Val))
	}

	// AccessPaths
	prop, err = oleutil.GetProperty(p.handle, "AccessPaths")
	if err != nil {
		return fmt.Errorf("oleutil.GetProperty(AccessPaths): %w", err)
	}
	for _, pa := range prop.ToArray().ToValueArray() {
		conv, ok := pa.(string)
		if !ok {
			return errors.New("error converting access path")
		}
		p.AccessPaths = append(p.AccessPaths, conv)
	}

	// GptType
	prop, err = oleutil.GetProperty(p.handle, "GptType")
	if err != nil {
		return fmt.Errorf("oleutil.GetProperty(GptType): %w", err)
	}
	p.GptType = prop.ToString()

	// GUID
	prop, err = oleutil.GetProperty(p.handle, "Guid")
	if err != nil {
		return fmt.Errorf("oleutil.GetProperty(Guid): %w", err)
	}
	p.GUID = prop.ToString()

	// All the non-strings
	for _, prop := range [][]interface{}{
		[]interface{}{"DiskNumber", &p.DiskNumber},
		[]interface{}{"PartitionNumber", &p.PartitionNumber},
		[]interface{}{"OperationalStatus", &p.OperationalStatus},
		[]interface{}{"TransitionState", &p.TransitionState},
		[]interface{}{"Size", &p.Size},
		[]interface{}{"MbrType", &p.MbrType},
		[]interface{}{"IsReadOnly", &p.IsReadOnly},
		[]interface{}{"IsOffline", &p.IsOffline},
		[]interface{}{"IsSystem", &p.IsSystem},
		[]interface{}{"IsBoot", &p.IsBoot},
		[]interface{}{"IsActive", &p.IsActive},
		[]interface{}{"IsHidden", &p.IsHidden},
		[]interface{}{"IsShadowCopy", &p.IsShadowCopy},
		[]interface{}{"NoDefaultDriveLetter", &p.NoDefaultDriveLetter},
	} {
		val, err := oleutil.GetProperty(p.handle, prop[0].(string))
		if err != nil {
			return fmt.Errorf("oleutil.GetProperty(%s): %w", prop[0].(string), err)
		}
		if err := assignVariant(val.Value(), prop[1]); err != nil {
			logger.Warningf("assignVariant(%s): %v", prop[0].(string), err)
		}
	}
	return nil
}

// Resize attempts to resize a partition.
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-resize
func (p *Partition) Resize(size uint64) (ExtendedStatus, error) {
	stat := ExtendedStatus{}
	var extendedStatus ole.VARIANT
	ole.VariantInit(&extendedStatus)
	// Convert the unint to a string because of because of https://docs.microsoft.com/en-us/previous-versions//aa393262%28v=vs.85%29?redirectedfrom=MSDN
	resultRaw, err := oleutil.CallMethod(p.handle, "Resize", strconv.FormatUint(size, 10), &extendedStatus)
	if err != nil {
		return stat, fmt.Errorf("Resize: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return stat, fmt.Errorf("error code returned during resize: %d", val)
	}
	return stat, nil
}

// A PartitionSet contains one or more Partitions.
type PartitionSet struct {
	Partitions []Partition
}

// Close releases all Partition handles inside a PartitionSet.
func (s *PartitionSet) Close() {
	for _, p := range s.Partitions {
		p.Close()
	}
}

// GetSupportedSize retrieves the minimum and maximum sizes that the partition can be resized to
//
// Ref: https://docs.microsoft.com/en-us/previous-versions/windows/desktop/stormgmt/msft-partition-getsupportedsizes
func (p *Partition) GetSupportedSize() (PartitionSupportedSize, ExtendedStatus, error) {
	size := PartitionSupportedSize{}
	stat := ExtendedStatus{}

	var sizemin ole.VARIANT
	ole.VariantInit(&sizemin)
	var sizemax ole.VARIANT
	ole.VariantInit(&sizemax)
	var extendedstatus ole.VARIANT
	ole.VariantInit(&extendedstatus)

	resultRaw, err := oleutil.CallMethod(p.handle, "GetSupportedSize", &sizemin, &sizemax, &extendedstatus)
	if err != nil {
		return size, stat, fmt.Errorf("GetSupportedSize: %w", err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return size, stat, fmt.Errorf("error code returned during GetSupportedSize: %d", val)
	}

	// Convert the results from an interface to uint64
	// Results are returned as strings because of https://docs.microsoft.com/en-us/previous-versions//aa393262%28v=vs.85%29?redirectedfrom=MSDN
	size.SizeMin, err = strconv.ParseUint(sizemin.Value().(string), 10, 64)
	if err != nil {
		return size, stat, fmt.Errorf("error attempting to parse sizemin: %w", err)
	}
	size.SizeMax, err = strconv.ParseUint(sizemax.Value().(string), 10, 64)
	if err != nil {
		return size, stat, fmt.Errorf("error attempting to parse sizemax: %w", err)
	}
	return size, stat, nil
}

// GetPartitions queries for local partitions.
//
// Close() must be called on the resulting PartitionSet to ensure all disks are released.
//
// Get all partitions:
//		svc.GetPartitions("")
//
// To get specific partitions, provide a valid WMI query filter string, for example:
//		svc.GetPartitions("WHERE DiskNumber=1")
func (svc *Service) GetPartitions(filter string) (PartitionSet, error) {
	parts := PartitionSet{}
	query := "SELECT * FROM MSFT_Partition"
	if filter != "" {
		query = fmt.Sprintf("%s %s", query, filter)
	}
	raw, err := oleutil.CallMethod(svc.wmiSvc, "ExecQuery", query)
	if err != nil {
		return parts, fmt.Errorf("ExecQuery(%s): %w", query, err)
	}
	result := raw.ToIDispatch()
	defer result.Release()

	countVar, err := oleutil.GetProperty(result, "Count")
	if err != nil {
		return parts, fmt.Errorf("oleutil.GetProperty(Count): %w", err)
	}
	count := int(countVar.Val)

	for i := 0; i < count; i++ {
		part := Partition{}
		itemRaw, err := oleutil.CallMethod(result, "ItemIndex", i)
		if err != nil {
			return parts, fmt.Errorf("oleutil.CallMethod(ItemIndex, %d): %w", i, err)
		}
		part.handle = itemRaw.ToIDispatch()

		if err := part.Query(); err != nil {
			return parts, err
		}

		parts.Partitions = append(parts.Partitions, part)
	}

	return parts, nil
}
<|endoftext|>


<commit_before>package goat

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync/atomic"
)

// HTTPConnHandler handles incoming HTTP (TCP) network connections
type HTTPConnHandler struct {
}

// Handle incoming HTTP connections and serve
func (h HTTPConnHandler) Handle(l net.Listener, httpDoneChan chan bool) {
	// Create shutdown function
	go func(l net.Listener, httpDoneChan chan bool) {
		// Wait for done signal
		Static.ShutdownChan <- <-Static.ShutdownChan

		// Close listener
		l.Close()
		httpDoneChan <- true
	}(l, httpDoneChan)

	// Log API configuration
	if Static.Config.API {
		log.Println("API functionality enabled")
	}

	// Set up HTTP routes for handling functions
	http.HandleFunc("/", parseHTTP)

	// Serve HTTP requests
	http.Serve(l, nil)
}

// Parse incoming HTTP connections before making tracker calls
func parseHTTP(w http.ResponseWriter, r *http.Request) {
	// Count incoming connections
	atomic.AddInt64(&Static.HTTP.Current, 1)
	atomic.AddInt64(&Static.HTTP.Total, 1)

	// Parse querystring
	querystring := r.URL.Query()

	// Flatten arrays into single values
	query := map[string]string{}
	for k, v := range querystring {
		query[k] = v[0]
	}

	// Check if IP was previously set
	if _, ok := query["ip"]; !ok {
		// If no IP set, detect and store it in query map
		query["ip"] = strings.Split(r.RemoteAddr, ":")[0]
	}

	// Add header to identify goat
	w.Header().Add("Server", fmt.Sprintf("%s/%s", App, Version))

	// Store current URL path
	url := r.URL.Path

	// Split URL into segments
	urlArr := strings.Split(url, "/")

	// If configured, Detect if client is making an API call
	url = urlArr[1]
	if Static.Config.API && url == "api" {
		// Log API calls
		log.Printf("API: %s\n", r.URL.Path)

		// Handle API calls, output JSON
		w.Header().Add("Content-Type", "application/json")
		APIRouter(w, r)
		return
	}

	// Detect if passkey present in URL
	var passkey string
	if len(urlArr) == 3 {
		passkey = urlArr[1]
		url = urlArr[2]
	}

	// Verify that torrent client is advertising its User-Agent, so we can use a whitelist
	if _, ok := r.Header["User-Agent"]; !ok {
		w.Write(HTTPTrackerError("Your client is not identifying itself"))
		return
	}

	client := r.Header["User-Agent"][0]

	// If configured, verify that torrent client is on whitelist
	if Static.Config.Whitelist {
		whitelist := new(WhitelistRecord).Load(client, "client")
		if whitelist == (WhitelistRecord{}) || !whitelist.Approved {
			w.Write(HTTPTrackerError("Your client is not whitelisted"))

			// Block things like browsers and web crawlers, because they will just clutter up the table
			if strings.Contains(client, "Mozilla") || strings.Contains(client, "Opera") {
				return
			}

			// Insert unknown clients into list for later approval
			if whitelist == (WhitelistRecord{}) {
				whitelist.Client = client
				whitelist.Approved = false

				log.Printf("whitelist: detected new client '%s', awaiting manual approval", client)

				go whitelist.Save()
			}

			return
		}
	}

	// Put client in query map
	query["client"] = client

	// Check if server is configured for passkey announce
	if Static.Config.Passkey && passkey == "" {
		w.Write(HTTPTrackerError("No passkey found in announce URL"))
		return
	}

	// Validate passkey if needed
	user := new(UserRecord).Load(passkey, "passkey")
	if Static.Config.Passkey && user == (UserRecord{}) {
		w.Write(HTTPTrackerError("Invalid passkey"))
		return
	}

	// Put passkey in query map
	query["passkey"] = user.Passkey

	// Mark client as HTTP
	query["udp"] = "0"

	// Create channel to return response to client
	resChan := make(chan []byte)

	// Handle tracker functions via different URLs
	switch url {
	// Tracker announce
	case "announce":
		// Validate required parameter input
		required := []string{"info_hash", "ip", "port", "uploaded", "downloaded", "left"}
		// Validate required integer input
		reqInt := []string{"port", "uploaded", "downloaded", "left"}

		// Check for required parameters
		for _, r := range required {
			if _, ok := query[r]; !ok {
				w.Write(HTTPTrackerError("Missing required parameter: " + r))
				close(resChan)
				return
			}
		}

		// Check for all valid integers
		for _, r := range reqInt {
			if _, ok := query[r]; ok {
				_, err := strconv.Atoi(query[r])
				if err != nil {
					w.Write(HTTPTrackerError("Invalid integer parameter: " + r))
					close(resChan)
					return
				}
			}
		}

		// Only allow compact announce
		if _, ok := query["compact"]; !ok || query["compact"] != "1" {
			w.Write(HTTPTrackerError("Your client does not support compact announce"))
			close(resChan)
			return
		}

		// Perform tracker announce
		go TrackerAnnounce(user, query, nil, resChan)
	// Tracker scrape
	case "scrape":
		go TrackerScrape(user, query, resChan)
	// Any undefined handlers
	default:
		w.Write(HTTPTrackerError("Malformed announce"))
		close(resChan)
		return
	}

	// Wait for response, and send it when ready
	w.Write(<-resChan)
	close(resChan)
	return
}
<commit_msg>Add API disabled error to API output<commit_after>package goat

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync/atomic"
)

// HTTPConnHandler handles incoming HTTP (TCP) network connections
type HTTPConnHandler struct {
}

// Handle incoming HTTP connections and serve
func (h HTTPConnHandler) Handle(l net.Listener, httpDoneChan chan bool) {
	// Create shutdown function
	go func(l net.Listener, httpDoneChan chan bool) {
		// Wait for done signal
		Static.ShutdownChan <- <-Static.ShutdownChan

		// Close listener
		l.Close()
		httpDoneChan <- true
	}(l, httpDoneChan)

	// Log API configuration
	if Static.Config.API {
		log.Println("API functionality enabled")
	}

	// Set up HTTP routes for handling functions
	http.HandleFunc("/", parseHTTP)

	// Serve HTTP requests
	http.Serve(l, nil)
}

// Parse incoming HTTP connections before making tracker calls
func parseHTTP(w http.ResponseWriter, r *http.Request) {
	// Count incoming connections
	atomic.AddInt64(&Static.HTTP.Current, 1)
	atomic.AddInt64(&Static.HTTP.Total, 1)

	// Parse querystring
	querystring := r.URL.Query()

	// Flatten arrays into single values
	query := map[string]string{}
	for k, v := range querystring {
		query[k] = v[0]
	}

	// Check if IP was previously set
	if _, ok := query["ip"]; !ok {
		// If no IP set, detect and store it in query map
		query["ip"] = strings.Split(r.RemoteAddr, ":")[0]
	}

	// Add header to identify goat
	w.Header().Add("Server", fmt.Sprintf("%s/%s", App, Version))

	// Store current URL path
	url := r.URL.Path

	// Split URL into segments
	urlArr := strings.Split(url, "/")

	// If configured, Detect if client is making an API call
	url = urlArr[1]
	if url == "api" {
		// Output JSON
		w.Header().Add("Content-Type", "application/json")

		// API enabled
		if Static.Config.API {
			// Log API calls
			log.Printf("API: %s\n", r.URL.Path)

			// Handle API calls, output JSON
			APIRouter(w, r)
			return
		} else {
			http.Error(w, string(APIErrorResponse("API is currently disabled")), 503)
			return
		}
	}

	// Detect if passkey present in URL
	var passkey string
	if len(urlArr) == 3 {
		passkey = urlArr[1]
		url = urlArr[2]
	}

	// Verify that torrent client is advertising its User-Agent, so we can use a whitelist
	if _, ok := r.Header["User-Agent"]; !ok {
		w.Write(HTTPTrackerError("Your client is not identifying itself"))
		return
	}

	client := r.Header["User-Agent"][0]

	// If configured, verify that torrent client is on whitelist
	if Static.Config.Whitelist {
		whitelist := new(WhitelistRecord).Load(client, "client")
		if whitelist == (WhitelistRecord{}) || !whitelist.Approved {
			w.Write(HTTPTrackerError("Your client is not whitelisted"))

			// Block things like browsers and web crawlers, because they will just clutter up the table
			if strings.Contains(client, "Mozilla") || strings.Contains(client, "Opera") {
				return
			}

			// Insert unknown clients into list for later approval
			if whitelist == (WhitelistRecord{}) {
				whitelist.Client = client
				whitelist.Approved = false

				log.Printf("whitelist: detected new client '%s', awaiting manual approval", client)

				go whitelist.Save()
			}

			return
		}
	}

	// Put client in query map
	query["client"] = client

	// Check if server is configured for passkey announce
	if Static.Config.Passkey && passkey == "" {
		w.Write(HTTPTrackerError("No passkey found in announce URL"))
		return
	}

	// Validate passkey if needed
	user := new(UserRecord).Load(passkey, "passkey")
	if Static.Config.Passkey && user == (UserRecord{}) {
		w.Write(HTTPTrackerError("Invalid passkey"))
		return
	}

	// Put passkey in query map
	query["passkey"] = user.Passkey

	// Mark client as HTTP
	query["udp"] = "0"

	// Create channel to return response to client
	resChan := make(chan []byte)

	// Handle tracker functions via different URLs
	switch url {
	// Tracker announce
	case "announce":
		// Validate required parameter input
		required := []string{"info_hash", "ip", "port", "uploaded", "downloaded", "left"}
		// Validate required integer input
		reqInt := []string{"port", "uploaded", "downloaded", "left"}

		// Check for required parameters
		for _, r := range required {
			if _, ok := query[r]; !ok {
				w.Write(HTTPTrackerError("Missing required parameter: " + r))
				close(resChan)
				return
			}
		}

		// Check for all valid integers
		for _, r := range reqInt {
			if _, ok := query[r]; ok {
				_, err := strconv.Atoi(query[r])
				if err != nil {
					w.Write(HTTPTrackerError("Invalid integer parameter: " + r))
					close(resChan)
					return
				}
			}
		}

		// Only allow compact announce
		if _, ok := query["compact"]; !ok || query["compact"] != "1" {
			w.Write(HTTPTrackerError("Your client does not support compact announce"))
			close(resChan)
			return
		}

		// Perform tracker announce
		go TrackerAnnounce(user, query, nil, resChan)
	// Tracker scrape
	case "scrape":
		go TrackerScrape(user, query, resChan)
	// Any undefined handlers
	default:
		w.Write(HTTPTrackerError("Malformed announce"))
		close(resChan)
		return
	}

	// Wait for response, and send it when ready
	w.Write(<-resChan)
	close(resChan)
	return
}
<|endoftext|>


<commit_before>package nicehash

type GlobalStats struct {
	Algo                  AlgoType `json:"algo"`
	ProfitabilityAboveBtc float32  `json:"profitability_above_btc,string"`
	ProfitabilityAboveLtc float32  `json:"profitability_above_ltc,string"`
	Price                 float64  `json:"price,string"`
	ProfitabilityBtc      float32  `json:"profitability_btc,string"`
	ProfitabilityLtc      float32  `json:"profitability_ltc,string"`
	Speed                 float64  `json:"speed,string"`
}

func (client *NicehashClient) GetStatsGlobalCurrent() ([]GlobalStats, error) {
	stats := &struct {
		Result struct {
			       Stats []GlobalStats `json:"stats"`
		       } `json:"result"`
	}{}
	params := &Params{Method:"stats.global.current", Algo:AlgoTypeMAX, Location:LocationMAX}
	_, err := client.sling.New().Get("").QueryStruct(params).ReceiveSuccess(&stats)
	if err != nil {
		return stats.Result.Stats, err
	}

	return stats.Result.Stats, nil
}

func (client *NicehashClient) GetStatsGlobalDay() ([]GlobalStats, error) {
	stats := &struct {
		Result struct {
			       Stats []GlobalStats `json:"stats"`
		       } `json:"result"`
	}{}
	params := &Params{Method:"stats.global.24h", Algo:AlgoTypeMAX, Location:LocationMAX}
	_, err := client.sling.New().Get("").QueryStruct(params).ReceiveSuccess(&stats)
	if err != nil {
		return stats.Result.Stats, err
	}

	return stats.Result.Stats, nil
}

type ProviderStats struct {
	Algo          AlgoType `json:"algo"`
	Balance       float64  `json:"balance,string"`
	AcceptedSpeed float64  `json:"accepted_speed,string"`
	RejectedSpeed float64  `json:"rejected_speed,string"`
}

func (client *NicehashClient) GetStatsProvider(addr string) ([]ProviderStats, error) {
	stats := &struct {
		Result struct {
			       Stats []ProviderStats `json:"stats"`
		       } `json:"result"`
	}{}
	params := &Params{Method:"stats.provider", Algo:AlgoTypeMAX, Location:LocationMAX, Addr:addr}
	_, err := client.sling.New().Get("").QueryStruct(params).ReceiveSuccess(&stats)
	if err != nil {
		return stats.Result.Stats, err
	}

	return stats.Result.Stats, nil
}

<commit_msg>Reformat code<commit_after>package nicehash

import (
	"errors"
)

type GlobalStats struct {
	Algo                  AlgoType `json:"algo"`
	ProfitabilityAboveBtc float32  `json:"profitability_above_btc,string"`
	ProfitabilityAboveLtc float32  `json:"profitability_above_ltc,string"`
	Price                 float64  `json:"price,string"`
	ProfitabilityBtc      float32  `json:"profitability_btc,string"`
	ProfitabilityLtc      float32  `json:"profitability_ltc,string"`
	Speed                 float64  `json:"speed,string"`
}

func (client *NicehashClient) GetStatsGlobalCurrent() ([]GlobalStats, error) {
	stats := &struct {
		Result struct {
			Error string        `json:"error"`
			Stats []GlobalStats `json:"stats"`
		} `json:"result"`
	}{}
	params := &Params{Method: "stats.global.current", Algo: AlgoTypeMAX, Location: LocationMAX}
	_, err := client.sling.New().Get("").QueryStruct(params).ReceiveSuccess(&stats)
	if err != nil {
		return nil, err
	}
	if stats.Result.Error != "" {
		return nil, errors.New(stats.Result.Error)
	}
	return stats.Result.Stats, nil
}

func (client *NicehashClient) GetStatsGlobalDay() ([]GlobalStats, error) {
	stats := &struct {
		Result struct {
			Error string        `json:"error"`
			Stats []GlobalStats `json:"stats"`
		} `json:"result"`
	}{}
	params := &Params{Method: "stats.global.24h", Algo: AlgoTypeMAX, Location: LocationMAX}
	_, err := client.sling.New().Get("").QueryStruct(params).ReceiveSuccess(&stats)
	if err != nil {
		return nil, err
	}
	if stats.Result.Error != "" {
		return nil, errors.New(stats.Result.Error)
	}
	return stats.Result.Stats, nil
}

type ProviderStats struct {
	Algo          AlgoType `json:"algo"`
	Balance       float64  `json:"balance,string"`
	AcceptedSpeed float64  `json:"accepted_speed,string"`
	RejectedSpeed float64  `json:"rejected_speed,string"`
}

func (client *NicehashClient) GetStatsProvider(addr string) ([]ProviderStats, error) {
	stats := &struct {
		Result struct {
			Error    string             `json:"error"`
			Stats    []ProviderStats    `json:"stats"`
		} `json:"result"`
	}{}
	params := &Params{Method: "stats.provider", Algo: AlgoTypeMAX, Location: LocationMAX, Addr: addr}
	_, err := client.sling.New().Get("").QueryStruct(params).ReceiveSuccess(&stats)
	if err != nil {
		return nil, err
	}
	if stats.Result.Error != "" {
		return nil, errors.New(stats.Result.Error)
	}
	return stats.Result.Stats, nil
}

<|endoftext|>


<commit_before>package simple

import (
	"github.com/loadimpact/speedboat/runner"
	"golang.org/x/net/context"
	"net/http"
	"runtime"
	"time"
)

type SimpleRunner struct {
	URL    string
	Client *http.Client
}

func New() *SimpleRunner {
	return &SimpleRunner{
		Client: &http.Client{
			Transport: &http.Transport{
				DisableKeepAlives: true,
			},
		},
	}
}

func (r *SimpleRunner) Run(ctx context.Context) <-chan runner.Result {
	ch := make(chan runner.Result)

	go func() {
		defer close(ch)

		for {
			// Note that we abort if we cannot create a request. This means we're either out of
			// memory, or we have invalid user input, neither of which are recoverable.
			req, err := http.NewRequest("GET", r.URL, nil)
			if err != nil {
				ch <- runner.Result{Error: err}
				return
			}
			req.Close = true
			req.Cancel = ctx.Done()

			mem := runtime.MemStats{}
			runtime.ReadMemStats(&mem)
			oldNumGC := mem.NumGC

			startTime := time.Now()
			res, err := r.Client.Do(req)
			duration := time.Since(startTime)

			// Accounting for one less GC pause than what's actually reported, otherwise we seem
			// to be getting timings with negative times taken... There needs to be a better way.
			runtime.ReadMemStats(&mem)
			for i := oldNumGC; i < mem.NumGC-1; i++ {
				duration -= time.Duration(mem.PauseNs[(i+255)%256])
			}

			select {
			case <-ctx.Done():
				return
			default:
				if err != nil {
					ch <- runner.Result{Error: err, Time: duration}
					continue
				}
				res.Body.Close()
				ch <- runner.Result{Time: duration}
			}
		}
	}()

	return ch
}
<commit_msg>ReadMemStats appears to have a mutex :/<commit_after>package simple

import (
	"github.com/loadimpact/speedboat/runner"
	"golang.org/x/net/context"
	"net/http"
	"time"
)

type SimpleRunner struct {
	URL    string
	Client *http.Client
}

func New() *SimpleRunner {
	return &SimpleRunner{
		Client: &http.Client{
			Transport: &http.Transport{
				DisableKeepAlives: true,
			},
		},
	}
}

func (r *SimpleRunner) Run(ctx context.Context) <-chan runner.Result {
	ch := make(chan runner.Result)

	go func() {
		defer close(ch)

		for {
			// Note that we abort if we cannot create a request. This means we're either out of
			// memory, or we have invalid user input, neither of which are recoverable.
			req, err := http.NewRequest("GET", r.URL, nil)
			if err != nil {
				ch <- runner.Result{Error: err}
				return
			}
			req.Close = true
			req.Cancel = ctx.Done()

			startTime := time.Now()
			res, err := r.Client.Do(req)
			duration := time.Since(startTime)

			select {
			case <-ctx.Done():
				return
			default:
				if err != nil {
					ch <- runner.Result{Error: err, Time: duration}
					continue
				}
				res.Body.Close()
				ch <- runner.Result{Time: duration}
			}
		}
	}()

	return ch
}
<|endoftext|>


<commit_before>package file

import (
	"github.com/watermint/toolbox/domain/dropbox/api/dbx_context"
	"github.com/watermint/toolbox/domain/dropbox/api/dbx_context_impl"
	"github.com/watermint/toolbox/domain/dropbox/filesystem"
	"github.com/watermint/toolbox/domain/dropbox/filesystem/dfs_dbx_to_local"
	"github.com/watermint/toolbox/domain/dropbox/filesystem/dfs_dbx_to_local_block"
	"github.com/watermint/toolbox/domain/dropbox/model/mo_file"
	"github.com/watermint/toolbox/domain/dropbox/model/mo_path"
	"github.com/watermint/toolbox/essentials/file/es_filesystem"
	"github.com/watermint/toolbox/essentials/file/es_filesystem_local"
	"github.com/watermint/toolbox/essentials/file/es_sync"
	"github.com/watermint/toolbox/essentials/log/esl"
	"github.com/watermint/toolbox/essentials/model/mo_filter"
	mo_path2 "github.com/watermint/toolbox/essentials/model/mo_path"
	"github.com/watermint/toolbox/infra/app"
	"github.com/watermint/toolbox/infra/control/app_control"
	"github.com/watermint/toolbox/infra/recipe/rc_exec"
	"github.com/watermint/toolbox/infra/recipe/rc_recipe"
	"github.com/watermint/toolbox/infra/report/rp_model"
	"github.com/watermint/toolbox/infra/ui/app_msg"
	"github.com/watermint/toolbox/quality/recipe/qtr_endtoend"
	"path/filepath"
)

type Download struct {
	Context     dbx_context.Context
	Delete      bool
	Overwrite   bool
	LocalPath   mo_path2.FileSystemPath
	DropboxPath mo_path.DropboxPath
	Downloaded  rp_model.TransactionReport
	Skipped     rp_model.TransactionReport
	Deleted     rp_model.RowReport
	Summary     rp_model.RowReport
	Name        mo_filter.Filter
}

func (z *Download) Preset() {
	z.Downloaded.SetModel(&mo_file.ConcreteEntry{}, &es_filesystem.EntryData{}, rp_model.HiddenColumns(
		"result.file_system_type",
		"result.name",
		"result.size",
		"result.mod_time",
		"result.is_file",
		"result.is_folder",
		"input.id",
		"input.tag",
		"input.path_lower",
		"input.revision",
		"input.shared_folder_id",
		"input.parent_shared_folder_id",
	))
	z.Skipped.SetModel(&es_filesystem.PathData{}, nil, rp_model.HiddenColumns(
		"input.file_system_type",
		"input.attributes",
		"input.entry_namespace.file_system_type",
		"input.entry_namespace.namespace_id",
		"input.entry_namespace.attributes",
	))
	z.Deleted.SetModel(&es_filesystem.PathData{}, rp_model.HiddenColumns(
		"file_system_type",
		"attributes",
		"entry_namespace.file_system_type",
		"entry_namespace.namespace_id",
		"entry_namespace.attributes",
	))
	z.Summary.SetModel(&Summary{})
}

func (z *Download) Exec(c app_control.Control) error {
	l := c.Log().With(esl.String("src", z.LocalPath.Path()), esl.String("dest", z.DropboxPath.Path()))
	localPath, err := filepath.Abs(z.LocalPath.Path())
	if err != nil {
		l.Debug("Unable to calc abs path", esl.Error(err), esl.String("localPath", z.LocalPath.Path()))
		return err
	}
	localPath = filepath.Clean(localPath)
	if err := z.Downloaded.Open(rp_model.NoConsoleOutput()); err != nil {
		return err
	}
	if err := z.Skipped.Open(rp_model.NoConsoleOutput()); err != nil {
		return err
	}
	if err := z.Deleted.Open(rp_model.NoConsoleOutput()); err != nil {
		return err
	}
	if err := z.Summary.Open(); err != nil {
		return err
	}
	l.Debug("Start downloading")

	srcFs := filesystem.NewFileSystem(z.Context)
	tgtFs := es_filesystem_local.NewFileSystem()
	var conn es_filesystem.Connector
	if c.Feature().Experiment(app.ExperimentDbxDownloadBlock) {
		l.Debug("Use range request copier")
		conn = dfs_dbx_to_local_block.NewDropboxToLocal(z.Context)
	} else {
		conn = dfs_dbx_to_local.NewDropboxToLocal(z.Context)
	}

	mustToDbxEntry := func(entry es_filesystem.Entry) mo_file.Entry {
		e, errConvert := filesystem.ToDropboxEntry(entry)
		if errConvert != nil {
			l.Debug("Unable ot convert", esl.Error(errConvert))
			panic("internal error")
		}
		return e
	}

	status := &Status{}
	status.start()

	syncer := es_sync.New(
		c.Log(),
		c.NewQueue(),
		srcFs,
		tgtFs,
		conn,
		es_sync.SyncDelete(z.Delete),
		es_sync.SyncOverwrite(z.Overwrite),
		es_sync.OnDeleteSuccess(func(target es_filesystem.Path) {
			status.delete()
			z.Deleted.Row(target.AsData())
		}),
		es_sync.OnDeleteFailure(func(target es_filesystem.Path, err es_filesystem.FileSystemError) {
			status.error()
		}),
		es_sync.OnCreateFolderSuccess(func(target es_filesystem.Path) {
			status.createFolder()
		}),
		es_sync.OnCreateFolderFailure(func(target es_filesystem.Path, err es_filesystem.FileSystemError) {
			status.error()
		}),
		es_sync.OnCopySuccess(func(source es_filesystem.Entry, target es_filesystem.Entry) {
			z.Downloaded.Success(mustToDbxEntry(source).Concrete(), target.AsData())
			status.download(source.Size())
		}),
		es_sync.OnCopyFailure(func(source es_filesystem.Path, err es_filesystem.FileSystemError) {
			status.error()
		}),
		es_sync.OnSkip(func(reason es_sync.SkipReason, source es_filesystem.Entry, target es_filesystem.Path) {
			var reasonMsg app_msg.Message
			switch reason {
			case es_sync.SkipExists:
				reasonMsg = MUpload.SkipExists
			case es_sync.SkipFilter:
				reasonMsg = MUpload.SkipFilter
			case es_sync.SkipSame:
				reasonMsg = MUpload.SkipSame
			default:
				reasonMsg = MUpload.SkipOther.With("Reason", reason)
			}
			z.Skipped.Skip(reasonMsg, source.Path().AsData())
			status.skip()
		}),
		es_sync.WithNameFilter(z.Name),
	)

	syncErr := syncer.Sync(filesystem.NewPath("", z.DropboxPath), es_filesystem_local.NewPath(z.LocalPath.Path()))
	if syncErr != nil {
		l.Debug("Sync finished with an error", esl.Error(syncErr))
	}
	status.finish()
	z.Summary.Row(status.summary)
	return syncErr
}

func (z *Download) Test(c app_control.Control) error {
	return rc_exec.ExecMock(c, &Download{}, func(r rc_recipe.Recipe) {
		m := r.(*Download)
		m.Context = dbx_context_impl.NewMock("mock", c)
		m.LocalPath = qtr_endtoend.NewTestFileSystemFolderPath(c, "down")
		m.DropboxPath = qtr_endtoend.NewTestDropboxFolderPath("down")
	})
}
<commit_msg>#637 : add log<commit_after>package file

import (
	"github.com/watermint/toolbox/domain/dropbox/api/dbx_context"
	"github.com/watermint/toolbox/domain/dropbox/api/dbx_context_impl"
	"github.com/watermint/toolbox/domain/dropbox/filesystem"
	"github.com/watermint/toolbox/domain/dropbox/filesystem/dfs_dbx_to_local"
	"github.com/watermint/toolbox/domain/dropbox/filesystem/dfs_dbx_to_local_block"
	"github.com/watermint/toolbox/domain/dropbox/model/mo_file"
	"github.com/watermint/toolbox/domain/dropbox/model/mo_path"
	"github.com/watermint/toolbox/essentials/file/es_filesystem"
	"github.com/watermint/toolbox/essentials/file/es_filesystem_local"
	"github.com/watermint/toolbox/essentials/file/es_sync"
	"github.com/watermint/toolbox/essentials/log/esl"
	"github.com/watermint/toolbox/essentials/model/mo_filter"
	mo_path2 "github.com/watermint/toolbox/essentials/model/mo_path"
	"github.com/watermint/toolbox/infra/app"
	"github.com/watermint/toolbox/infra/control/app_control"
	"github.com/watermint/toolbox/infra/recipe/rc_exec"
	"github.com/watermint/toolbox/infra/recipe/rc_recipe"
	"github.com/watermint/toolbox/infra/report/rp_model"
	"github.com/watermint/toolbox/infra/ui/app_msg"
	"github.com/watermint/toolbox/quality/recipe/qtr_endtoend"
	"path/filepath"
)

type Download struct {
	Context     dbx_context.Context
	Delete      bool
	Overwrite   bool
	LocalPath   mo_path2.FileSystemPath
	DropboxPath mo_path.DropboxPath
	Downloaded  rp_model.TransactionReport
	Skipped     rp_model.TransactionReport
	Deleted     rp_model.RowReport
	Summary     rp_model.RowReport
	Name        mo_filter.Filter
}

func (z *Download) Preset() {
	z.Downloaded.SetModel(&mo_file.ConcreteEntry{}, &es_filesystem.EntryData{}, rp_model.HiddenColumns(
		"result.file_system_type",
		"result.name",
		"result.size",
		"result.mod_time",
		"result.is_file",
		"result.is_folder",
		"input.id",
		"input.tag",
		"input.path_lower",
		"input.revision",
		"input.shared_folder_id",
		"input.parent_shared_folder_id",
	))
	z.Skipped.SetModel(&es_filesystem.PathData{}, nil, rp_model.HiddenColumns(
		"input.file_system_type",
		"input.attributes",
		"input.entry_namespace.file_system_type",
		"input.entry_namespace.namespace_id",
		"input.entry_namespace.attributes",
	))
	z.Deleted.SetModel(&es_filesystem.PathData{}, rp_model.HiddenColumns(
		"file_system_type",
		"attributes",
		"entry_namespace.file_system_type",
		"entry_namespace.namespace_id",
		"entry_namespace.attributes",
	))
	z.Summary.SetModel(&Summary{})
}

func (z *Download) Exec(c app_control.Control) error {
	l := c.Log().With(esl.String("src", z.LocalPath.Path()), esl.String("dest", z.DropboxPath.Path()))
	localPath, err := filepath.Abs(z.LocalPath.Path())
	if err != nil {
		l.Debug("Unable to calc abs path", esl.Error(err), esl.String("localPath", z.LocalPath.Path()))
		return err
	}
	localPath = filepath.Clean(localPath)
	if err := z.Downloaded.Open(rp_model.NoConsoleOutput()); err != nil {
		return err
	}
	if err := z.Skipped.Open(rp_model.NoConsoleOutput()); err != nil {
		return err
	}
	if err := z.Deleted.Open(rp_model.NoConsoleOutput()); err != nil {
		return err
	}
	if err := z.Summary.Open(); err != nil {
		return err
	}
	l.Debug("Start downloading")

	srcFs := filesystem.NewFileSystem(z.Context)
	tgtFs := es_filesystem_local.NewFileSystem()
	var conn es_filesystem.Connector
	if c.Feature().Experiment(app.ExperimentDbxDownloadBlock) {
		l.Debug("Use block copier")
		conn = dfs_dbx_to_local_block.NewDropboxToLocal(z.Context)
	} else {
		l.Debug("Use standard copier")
		conn = dfs_dbx_to_local.NewDropboxToLocal(z.Context)
	}

	mustToDbxEntry := func(entry es_filesystem.Entry) mo_file.Entry {
		e, errConvert := filesystem.ToDropboxEntry(entry)
		if errConvert != nil {
			l.Debug("Unable ot convert", esl.Error(errConvert))
			panic("internal error")
		}
		return e
	}

	status := &Status{}
	status.start()

	syncer := es_sync.New(
		c.Log(),
		c.NewQueue(),
		srcFs,
		tgtFs,
		conn,
		es_sync.SyncDelete(z.Delete),
		es_sync.SyncOverwrite(z.Overwrite),
		es_sync.OnDeleteSuccess(func(target es_filesystem.Path) {
			status.delete()
			z.Deleted.Row(target.AsData())
		}),
		es_sync.OnDeleteFailure(func(target es_filesystem.Path, err es_filesystem.FileSystemError) {
			status.error()
		}),
		es_sync.OnCreateFolderSuccess(func(target es_filesystem.Path) {
			status.createFolder()
		}),
		es_sync.OnCreateFolderFailure(func(target es_filesystem.Path, err es_filesystem.FileSystemError) {
			status.error()
		}),
		es_sync.OnCopySuccess(func(source es_filesystem.Entry, target es_filesystem.Entry) {
			z.Downloaded.Success(mustToDbxEntry(source).Concrete(), target.AsData())
			status.download(source.Size())
		}),
		es_sync.OnCopyFailure(func(source es_filesystem.Path, err es_filesystem.FileSystemError) {
			status.error()
		}),
		es_sync.OnSkip(func(reason es_sync.SkipReason, source es_filesystem.Entry, target es_filesystem.Path) {
			var reasonMsg app_msg.Message
			switch reason {
			case es_sync.SkipExists:
				reasonMsg = MUpload.SkipExists
			case es_sync.SkipFilter:
				reasonMsg = MUpload.SkipFilter
			case es_sync.SkipSame:
				reasonMsg = MUpload.SkipSame
			default:
				reasonMsg = MUpload.SkipOther.With("Reason", reason)
			}
			z.Skipped.Skip(reasonMsg, source.Path().AsData())
			status.skip()
		}),
		es_sync.WithNameFilter(z.Name),
	)

	syncErr := syncer.Sync(filesystem.NewPath("", z.DropboxPath), es_filesystem_local.NewPath(z.LocalPath.Path()))
	if syncErr != nil {
		l.Debug("Sync finished with an error", esl.Error(syncErr))
	}
	status.finish()
	z.Summary.Row(status.summary)
	return syncErr
}

func (z *Download) Test(c app_control.Control) error {
	return rc_exec.ExecMock(c, &Download{}, func(r rc_recipe.Recipe) {
		m := r.(*Download)
		m.Context = dbx_context_impl.NewMock("mock", c)
		m.LocalPath = qtr_endtoend.NewTestFileSystemFolderPath(c, "down")
		m.DropboxPath = qtr_endtoend.NewTestDropboxFolderPath("down")
	})
}
<|endoftext|>


<commit_before>package inputsocket

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/Sirupsen/logrus"

	"github.com/tsaikd/gogstash/config"
	"github.com/tsaikd/gogstash/config/logevent"
)

// ModuleName is the name used in config file
const ModuleName = "socket"

// InputConfig holds the output configuration json fields
type InputConfig struct {
	config.InputConfig
	Socket  string `json:"socket"`  // Type of socket, must be one of ["tcp", "unix", "unixpacket"].
	Address string `json:"address"` // For TCP, address must have the form `host:port`. For Unix networks, the address must be a file system path.
}

// DefaultInputConfig returns an InputConfig struct with default values
func DefaultInputConfig() InputConfig {
	return InputConfig{
		InputConfig: config.InputConfig{
			CommonConfig: config.CommonConfig{
				Type: ModuleName,
			},
		},
	}
}

// InitHandler initialize the input plugin
func InitHandler(confraw *config.ConfigRaw) (config.TypeInputConfig, error) {
	conf := DefaultInputConfig()
	if err := config.ReflectConfig(confraw, &conf); err != nil {
		return nil, err
	}
	return &conf, nil
}

// Start wraps the actual function starting the plugin
func (i *InputConfig) Start() {
	i.Invoke(i.start)
}

func (i *InputConfig) start(logger *logrus.Logger, evchan chan logevent.LogEvent) {

	switch i.Socket {
	case "unix", "unixpacket":
		// Remove existing unix socket
		os.Remove(i.Address)
	}

	l, err := net.Listen(i.Socket, i.Address)
	if err != nil {
		logger.Error(ModuleName, ": Unable to listen to socket.", err)
	}
	for {
		conn, err := l.Accept()
		if err != nil {
			logger.Error(ModuleName, ": socket accept error.", err)
		}
		go parse(conn, logger, evchan)
	}
}

func parse(conn net.Conn, logger *logrus.Logger, evchan chan logevent.LogEvent) {
	defer conn.Close()

	// Duplicate buffer to be able to read it even after failed json decoding
	buf, _ := ioutil.ReadAll(conn)
	js := bytes.NewBuffer(buf)
	raw := bytes.NewBuffer(buf)

	dec := json.NewDecoder(js)
	for {
		// Assume first the message is JSON and try to decode it
		var jsonMsg map[string]interface{}
		if err := dec.Decode(&jsonMsg); err == io.EOF {
			break
		} else if err != nil {
			// if decoding fail, use raw message in LogEvent "Message" field
			evchan <- logevent.LogEvent{
				Timestamp: time.Now(),
				Message:   raw.String(),
			}
			break
		}
		evchan <- logevent.LogEvent{Extra: jsonMsg}
	}
}
<commit_msg>Fix io.Reader duplication ioutil.ReadAll was blocking until the end of the connection<commit_after>package inputsocket

import (
	"bytes"
	"encoding/json"
	"io"
	"net"
	"os"
	"time"

	"github.com/Sirupsen/logrus"

	"github.com/tsaikd/gogstash/config"
	"github.com/tsaikd/gogstash/config/logevent"
)

// ModuleName is the name used in config file
const ModuleName = "socket"

// InputConfig holds the output configuration json fields
type InputConfig struct {
	config.InputConfig
	Socket  string `json:"socket"`  // Type of socket, must be one of ["tcp", "unix", "unixpacket"].
	Address string `json:"address"` // For TCP, address must have the form `host:port`. For Unix networks, the address must be a file system path.
}

// DefaultInputConfig returns an InputConfig struct with default values
func DefaultInputConfig() InputConfig {
	return InputConfig{
		InputConfig: config.InputConfig{
			CommonConfig: config.CommonConfig{
				Type: ModuleName,
			},
		},
	}
}

// InitHandler initialize the input plugin
func InitHandler(confraw *config.ConfigRaw) (config.TypeInputConfig, error) {
	conf := DefaultInputConfig()
	if err := config.ReflectConfig(confraw, &conf); err != nil {
		return nil, err
	}
	return &conf, nil
}

// Start wraps the actual function starting the plugin
func (i *InputConfig) Start() {
	i.Invoke(i.start)
}

func (i *InputConfig) start(logger *logrus.Logger, evchan chan logevent.LogEvent) {

	switch i.Socket {
	case "unix", "unixpacket":
		// Remove existing unix socket
		os.Remove(i.Address)
	}

	l, err := net.Listen(i.Socket, i.Address)
	if err != nil {
		logger.Error(ModuleName, ": Unable to listen to socket.", err)
	}
	for {
		conn, err := l.Accept()
		if err != nil {
			logger.Error(ModuleName, ": socket accept error.", err)
		}
		go parse(conn, logger, evchan)
	}
}

func parse(conn net.Conn, logger *logrus.Logger, evchan chan logevent.LogEvent) {
	defer conn.Close()

	// Duplicate buffer to be able to read it even after failed json decoding
	var streamCopy bytes.Buffer
	stream := io.TeeReader(conn, &streamCopy)

	dec := json.NewDecoder(stream)
	for {
		// Assume first the message is JSON and try to decode it
		var jsonMsg map[string]interface{}
		if err := dec.Decode(&jsonMsg); err == io.EOF {
			break
		} else if err != nil {
			// If decoding fail, split raw message by line
			// and send a log event per line
			for {
				line, err := streamCopy.ReadString('\n')
				evchan <- logevent.LogEvent{
					Timestamp: time.Now(),
					Message:   line,
				}
				if err != nil {
					break
				}
			}
			break
		}
		evchan <- logevent.LogEvent{Extra: jsonMsg}
	}
}
<|endoftext|>