<commit_before>// +build windows

package host

import (
	"context"
	"fmt"
	"math"
	"os"
	"runtime"
	"strings"
	"sync/atomic"
	"syscall"
	"time"
	"unsafe"

	"github.com/StackExchange/wmi"
	"github.com/shirou/gopsutil/internal/common"
	process "github.com/shirou/gopsutil/process"
	"golang.org/x/sys/windows"
)

var (
	procGetSystemTimeAsFileTime = common.Modkernel32.NewProc("GetSystemTimeAsFileTime")
	procGetTickCount32          = common.Modkernel32.NewProc("GetTickCount")
	procGetTickCount64          = common.Modkernel32.NewProc("GetTickCount64")
	procRtlGetVersion           = common.ModNt.NewProc("RtlGetVersion")
)

// https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/content/wdm/ns-wdm-_osversioninfoexw
type osVersionInfoExW struct {
	dwOSVersionInfoSize uint32
	dwMajorVersion      uint32
	dwMinorVersion      uint32
	dwBuildNumber       uint32
	dwPlatformId        uint32
	szCSDVersion        [128]uint16
	wServicePackMajor   uint16
	wServicePackMinor   uint16
	wSuiteMask          uint16
	wProductType        uint8
	wReserved           uint8
}

func Info() (*InfoStat, error) {
	return InfoWithContext(context.Background())
}

func InfoWithContext(ctx context.Context) (*InfoStat, error) {
	ret := &InfoStat{
		OS: runtime.GOOS,
	}

	{
		hostname, err := os.Hostname()
		if err == nil {
			ret.Hostname = hostname
		}
	}

	{
		platform, family, version, err := PlatformInformationWithContext(ctx)
		if err == nil {
			ret.Platform = platform
			ret.PlatformFamily = family
			ret.PlatformVersion = version
		} else {
			return ret, err
		}
	}

	{
		boot, err := BootTime()
		if err == nil {
			ret.BootTime = boot
			ret.Uptime, _ = Uptime()
		}
	}

	{
		hostID, err := getMachineGuid()
		if err == nil {
			ret.HostID = strings.ToLower(hostID)
		}
	}

	{
		procs, err := process.Pids()
		if err == nil {
			ret.Procs = uint64(len(procs))
		}
	}

	return ret, nil
}

func getMachineGuid() (string, error) {
	// there has been reports of issues on 32bit using golang.org/x/sys/windows/registry, see https://github.com/shirou/gopsutil/pull/312#issuecomment-277422612
	// for rationale of using windows.RegOpenKeyEx/RegQueryValueEx instead of registry.OpenKey/GetStringValue
	var h windows.Handle
	err := windows.RegOpenKeyEx(windows.HKEY_LOCAL_MACHINE, windows.StringToUTF16Ptr(`SOFTWARE\Microsoft\Cryptography`), 0, windows.KEY_READ|windows.KEY_WOW64_64KEY, &h)
	if err != nil {
		return "", err
	}
	defer windows.RegCloseKey(h)

	const windowsRegBufLen = 74 // len(`{`) + len(`abcdefgh-1234-456789012-123345456671` * 2) + len(`}`) // 2 == bytes/UTF16
	const uuidLen = 36

	var regBuf [windowsRegBufLen]uint16
	bufLen := uint32(windowsRegBufLen)
	var valType uint32
	err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`MachineGuid`), nil, &valType, (*byte)(unsafe.Pointer(&regBuf[0])), &bufLen)
	if err != nil {
		return "", err
	}

	hostID := windows.UTF16ToString(regBuf[:])
	hostIDLen := len(hostID)
	if hostIDLen != uuidLen {
		return "", fmt.Errorf("HostID incorrect: %q\n", hostID)
	}

	return hostID, nil
}

func Uptime() (uint64, error) {
	return UptimeWithContext(context.Background())
}

func UptimeWithContext(ctx context.Context) (uint64, error) {
	procGetTickCount := procGetTickCount64
	err := procGetTickCount64.Find()
	if err != nil {
		procGetTickCount = procGetTickCount32 // handle WinXP, but keep in mind that "the time will wrap around to zero if the system is run continuously for 49.7 days." from MSDN
	}
	r1, _, lastErr := syscall.Syscall(procGetTickCount.Addr(), 0, 0, 0, 0)
	if lastErr != 0 {
		return 0, lastErr
	}
	return uint64((time.Duration(r1) * time.Millisecond).Seconds()), nil
}

func bootTimeFromUptime(up uint64) uint64 {
	return uint64(time.Now().Unix()) - up
}

// cachedBootTime must be accessed via atomic.Load/StoreUint64
var cachedBootTime uint64

func BootTime() (uint64, error) {
	return BootTimeWithContext(context.Background())
}

func BootTimeWithContext(ctx context.Context) (uint64, error) {
	t := atomic.LoadUint64(&cachedBootTime)
	if t != 0 {
		return t, nil
	}
	up, err := Uptime()
	if err != nil {
		return 0, err
	}
	t = bootTimeFromUptime(up)
	atomic.StoreUint64(&cachedBootTime, t)
	return t, nil
}

func PlatformInformation() (platform string, family string, version string, err error) {
	return PlatformInformationWithContext(context.Background())
}

func PlatformInformationWithContext(ctx context.Context) (platform string, family string, version string, err error) {
	// GetVersionEx lies on Windows 8.1 and returns as Windows 8 if we don't declare compatibility in manifest
	// RtlGetVersion bypasses this lying layer and returns the true Windows version
	// https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/content/wdm/nf-wdm-rtlgetversion
	// https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/content/wdm/ns-wdm-_osversioninfoexw
	var osInfo osVersionInfoExW
	osInfo.dwOSVersionInfoSize = uint32(unsafe.Sizeof(osInfo))
	ret, _, err := procRtlGetVersion.Call(uintptr(unsafe.Pointer(&osInfo)))
	if ret != 0 {
		return
	}

	// Platform
	var h windows.Handle // like getMachineGuid(), we query the registry using the raw windows.RegOpenKeyEx/RegQueryValueEx
	err = windows.RegOpenKeyEx(windows.HKEY_LOCAL_MACHINE, windows.StringToUTF16Ptr(`SOFTWARE\Microsoft\Windows NT\CurrentVersion`), 0, windows.KEY_READ|windows.KEY_WOW64_64KEY, &h)
	if err != nil {
		return
	}
	defer windows.RegCloseKey(h)
	var bufLen uint32
	var valType uint32
	err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`ProductName`), nil, &valType, nil, &bufLen)
	if err != nil {
		return
	}
	regBuf := make([]uint16, bufLen/2+1)
	err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`ProductName`), nil, &valType, (*byte)(unsafe.Pointer(&regBuf[0])), &bufLen)
	if err != nil {
		return
	}
	platform = windows.UTF16ToString(regBuf[:])
	if !strings.HasPrefix(platform, "Microsoft") {
		platform = "Microsoft " + platform
	}
	err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`CSDVersion`), nil, &valType, nil, &bufLen) // append Service Pack number, only on success
	if err == nil {                                                                                // don't return an error if only the Service Pack retrieval fails
		regBuf = make([]uint16, bufLen/2+1)
		err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`CSDVersion`), nil, &valType, (*byte)(unsafe.Pointer(&regBuf[0])), &bufLen)
		if err == nil {
			platform += " " + windows.UTF16ToString(regBuf[:])
		}
	}

	// PlatformFamily
	switch osInfo.wProductType {
	case 1:
		family = "Standalone Workstation"
	case 2:
		family = "Server (Domain Controller)"
	case 3:
		family = "Server"
	}

	// Platform Version
	version = fmt.Sprintf("%d.%d.%d Build %d", osInfo.dwMajorVersion, osInfo.dwMinorVersion, osInfo.dwBuildNumber, osInfo.dwBuildNumber)

	return platform, family, version, nil
}

func Users() ([]UserStat, error) {
	return UsersWithContext(context.Background())
}

func UsersWithContext(ctx context.Context) ([]UserStat, error) {
	var ret []UserStat

	return ret, nil
}

func SensorsTemperatures() ([]TemperatureStat, error) {
	return SensorsTemperaturesWithContext(context.Background())
}

func SensorsTemperaturesWithContext(ctx context.Context) ([]TemperatureStat, error) {
	var ret []TemperatureStat
	var dst []MSAcpi_ThermalZoneTemperature
	q := wmi.CreateQuery(&dst, "")
	if err := common.WMIQueryWithContext(ctx, q, &dst, nil, "root/wmi"); err != nil {
		return ret, err
	}

	for _, v := range dst {

		ts := TemperatureStat{
			SensorKey:   v.InstanceName,
			Temperature: kelvinToCelsius(v.CurrentTemperature, 2),
		}
		ret = append(ret, ts)
	}

	return ret, nil
}

func kelvinToCelsius(temp uint32, n int) float64 {
	// wmi return temperature Kelvin * 10, so need to divide the result by 10,
	// and then minus 273.15 to get °Celsius.
	t := float64(temp/10) - 273.15
	n10 := math.Pow10(n)
	return math.Trunc((t+0.5/n10)*n10) / n10
}

func Virtualization() (string, string, error) {
	return VirtualizationWithContext(context.Background())
}

func VirtualizationWithContext(ctx context.Context) (string, string, error) {
	return "", "", common.ErrNotImplementedError
}

func KernelVersion() (string, error) {
	return KernelVersionWithContext(context.Background())
}

func KernelVersionWithContext(ctx context.Context) (string, error) {
	_, _, version, err := PlatformInformation()
	return version, err
}

type MSAcpi_ThermalZoneTemperature struct {
	Active             bool
	CriticalTripPoint  uint32
	CurrentTemperature uint32
	InstanceName       string
}
<commit_msg>[host][windows] Remove exported wmi type from #550<commit_after>// +build windows

package host

import (
	"context"
	"fmt"
	"math"
	"os"
	"runtime"
	"strings"
	"sync/atomic"
	"syscall"
	"time"
	"unsafe"

	"github.com/StackExchange/wmi"
	"github.com/shirou/gopsutil/internal/common"
	process "github.com/shirou/gopsutil/process"
	"golang.org/x/sys/windows"
)

var (
	procGetSystemTimeAsFileTime = common.Modkernel32.NewProc("GetSystemTimeAsFileTime")
	procGetTickCount32          = common.Modkernel32.NewProc("GetTickCount")
	procGetTickCount64          = common.Modkernel32.NewProc("GetTickCount64")
	procRtlGetVersion           = common.ModNt.NewProc("RtlGetVersion")
)

// https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/content/wdm/ns-wdm-_osversioninfoexw
type osVersionInfoExW struct {
	dwOSVersionInfoSize uint32
	dwMajorVersion      uint32
	dwMinorVersion      uint32
	dwBuildNumber       uint32
	dwPlatformId        uint32
	szCSDVersion        [128]uint16
	wServicePackMajor   uint16
	wServicePackMinor   uint16
	wSuiteMask          uint16
	wProductType        uint8
	wReserved           uint8
}

type msAcpi_ThermalZoneTemperature struct {
	Active             bool
	CriticalTripPoint  uint32
	CurrentTemperature uint32
	InstanceName       string
}

func Info() (*InfoStat, error) {
	return InfoWithContext(context.Background())
}

func InfoWithContext(ctx context.Context) (*InfoStat, error) {
	ret := &InfoStat{
		OS: runtime.GOOS,
	}

	{
		hostname, err := os.Hostname()
		if err == nil {
			ret.Hostname = hostname
		}
	}

	{
		platform, family, version, err := PlatformInformationWithContext(ctx)
		if err == nil {
			ret.Platform = platform
			ret.PlatformFamily = family
			ret.PlatformVersion = version
		} else {
			return ret, err
		}
	}

	{
		boot, err := BootTime()
		if err == nil {
			ret.BootTime = boot
			ret.Uptime, _ = Uptime()
		}
	}

	{
		hostID, err := getMachineGuid()
		if err == nil {
			ret.HostID = strings.ToLower(hostID)
		}
	}

	{
		procs, err := process.Pids()
		if err == nil {
			ret.Procs = uint64(len(procs))
		}
	}

	return ret, nil
}

func getMachineGuid() (string, error) {
	// there has been reports of issues on 32bit using golang.org/x/sys/windows/registry, see https://github.com/shirou/gopsutil/pull/312#issuecomment-277422612
	// for rationale of using windows.RegOpenKeyEx/RegQueryValueEx instead of registry.OpenKey/GetStringValue
	var h windows.Handle
	err := windows.RegOpenKeyEx(windows.HKEY_LOCAL_MACHINE, windows.StringToUTF16Ptr(`SOFTWARE\Microsoft\Cryptography`), 0, windows.KEY_READ|windows.KEY_WOW64_64KEY, &h)
	if err != nil {
		return "", err
	}
	defer windows.RegCloseKey(h)

	const windowsRegBufLen = 74 // len(`{`) + len(`abcdefgh-1234-456789012-123345456671` * 2) + len(`}`) // 2 == bytes/UTF16
	const uuidLen = 36

	var regBuf [windowsRegBufLen]uint16
	bufLen := uint32(windowsRegBufLen)
	var valType uint32
	err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`MachineGuid`), nil, &valType, (*byte)(unsafe.Pointer(&regBuf[0])), &bufLen)
	if err != nil {
		return "", err
	}

	hostID := windows.UTF16ToString(regBuf[:])
	hostIDLen := len(hostID)
	if hostIDLen != uuidLen {
		return "", fmt.Errorf("HostID incorrect: %q\n", hostID)
	}

	return hostID, nil
}

func Uptime() (uint64, error) {
	return UptimeWithContext(context.Background())
}

func UptimeWithContext(ctx context.Context) (uint64, error) {
	procGetTickCount := procGetTickCount64
	err := procGetTickCount64.Find()
	if err != nil {
		procGetTickCount = procGetTickCount32 // handle WinXP, but keep in mind that "the time will wrap around to zero if the system is run continuously for 49.7 days." from MSDN
	}
	r1, _, lastErr := syscall.Syscall(procGetTickCount.Addr(), 0, 0, 0, 0)
	if lastErr != 0 {
		return 0, lastErr
	}
	return uint64((time.Duration(r1) * time.Millisecond).Seconds()), nil
}

func bootTimeFromUptime(up uint64) uint64 {
	return uint64(time.Now().Unix()) - up
}

// cachedBootTime must be accessed via atomic.Load/StoreUint64
var cachedBootTime uint64

func BootTime() (uint64, error) {
	return BootTimeWithContext(context.Background())
}

func BootTimeWithContext(ctx context.Context) (uint64, error) {
	t := atomic.LoadUint64(&cachedBootTime)
	if t != 0 {
		return t, nil
	}
	up, err := Uptime()
	if err != nil {
		return 0, err
	}
	t = bootTimeFromUptime(up)
	atomic.StoreUint64(&cachedBootTime, t)
	return t, nil
}

func PlatformInformation() (platform string, family string, version string, err error) {
	return PlatformInformationWithContext(context.Background())
}

func PlatformInformationWithContext(ctx context.Context) (platform string, family string, version string, err error) {
	// GetVersionEx lies on Windows 8.1 and returns as Windows 8 if we don't declare compatibility in manifest
	// RtlGetVersion bypasses this lying layer and returns the true Windows version
	// https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/content/wdm/nf-wdm-rtlgetversion
	// https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/content/wdm/ns-wdm-_osversioninfoexw
	var osInfo osVersionInfoExW
	osInfo.dwOSVersionInfoSize = uint32(unsafe.Sizeof(osInfo))
	ret, _, err := procRtlGetVersion.Call(uintptr(unsafe.Pointer(&osInfo)))
	if ret != 0 {
		return
	}

	// Platform
	var h windows.Handle // like getMachineGuid(), we query the registry using the raw windows.RegOpenKeyEx/RegQueryValueEx
	err = windows.RegOpenKeyEx(windows.HKEY_LOCAL_MACHINE, windows.StringToUTF16Ptr(`SOFTWARE\Microsoft\Windows NT\CurrentVersion`), 0, windows.KEY_READ|windows.KEY_WOW64_64KEY, &h)
	if err != nil {
		return
	}
	defer windows.RegCloseKey(h)
	var bufLen uint32
	var valType uint32
	err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`ProductName`), nil, &valType, nil, &bufLen)
	if err != nil {
		return
	}
	regBuf := make([]uint16, bufLen/2+1)
	err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`ProductName`), nil, &valType, (*byte)(unsafe.Pointer(&regBuf[0])), &bufLen)
	if err != nil {
		return
	}
	platform = windows.UTF16ToString(regBuf[:])
	if !strings.HasPrefix(platform, "Microsoft") {
		platform = "Microsoft " + platform
	}
	err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`CSDVersion`), nil, &valType, nil, &bufLen) // append Service Pack number, only on success
	if err == nil {                                                                                // don't return an error if only the Service Pack retrieval fails
		regBuf = make([]uint16, bufLen/2+1)
		err = windows.RegQueryValueEx(h, windows.StringToUTF16Ptr(`CSDVersion`), nil, &valType, (*byte)(unsafe.Pointer(&regBuf[0])), &bufLen)
		if err == nil {
			platform += " " + windows.UTF16ToString(regBuf[:])
		}
	}

	// PlatformFamily
	switch osInfo.wProductType {
	case 1:
		family = "Standalone Workstation"
	case 2:
		family = "Server (Domain Controller)"
	case 3:
		family = "Server"
	}

	// Platform Version
	version = fmt.Sprintf("%d.%d.%d Build %d", osInfo.dwMajorVersion, osInfo.dwMinorVersion, osInfo.dwBuildNumber, osInfo.dwBuildNumber)

	return platform, family, version, nil
}

func Users() ([]UserStat, error) {
	return UsersWithContext(context.Background())
}

func UsersWithContext(ctx context.Context) ([]UserStat, error) {
	var ret []UserStat

	return ret, nil
}

func SensorsTemperatures() ([]TemperatureStat, error) {
	return SensorsTemperaturesWithContext(context.Background())
}

func SensorsTemperaturesWithContext(ctx context.Context) ([]TemperatureStat, error) {
	var ret []TemperatureStat
	var dst []msAcpi_ThermalZoneTemperature
	q := wmi.CreateQuery(&dst, "")
	if err := common.WMIQueryWithContext(ctx, q, &dst, nil, "root/wmi"); err != nil {
		return ret, err
	}

	for _, v := range dst {
		ts := TemperatureStat{
			SensorKey:   v.InstanceName,
			Temperature: kelvinToCelsius(v.CurrentTemperature, 2),
		}
		ret = append(ret, ts)
	}

	return ret, nil
}

func kelvinToCelsius(temp uint32, n int) float64 {
	// wmi return temperature Kelvin * 10, so need to divide the result by 10,
	// and then minus 273.15 to get °Celsius.
	t := float64(temp/10) - 273.15
	n10 := math.Pow10(n)
	return math.Trunc((t+0.5/n10)*n10) / n10
}

func Virtualization() (string, string, error) {
	return VirtualizationWithContext(context.Background())
}

func VirtualizationWithContext(ctx context.Context) (string, string, error) {
	return "", "", common.ErrNotImplementedError
}

func KernelVersion() (string, error) {
	return KernelVersionWithContext(context.Background())
}

func KernelVersionWithContext(ctx context.Context) (string, error) {
	_, _, version, err := PlatformInformation()
	return version, err
}
<|endoftext|>


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

/*
Package surfacers is the base package for creating Surfacer objects that are
used for writing metics data to different monitoring services.

Any Surfacer that is created for writing metrics data to a monitor system
should implement the below Surfacer interface and should accept
metrics.EventMetrics object through a Write() call. Each new surfacer should
also plug itself in through the New() method defined here.
*/
package surfacers

import (
	"context"
	"fmt"
	"html/template"
	"strings"
	"sync"

	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/metrics"
	"github.com/google/cloudprober/surfacers/file"
	"github.com/google/cloudprober/surfacers/postgres"
	"github.com/google/cloudprober/surfacers/prometheus"
	"github.com/google/cloudprober/surfacers/stackdriver"
	"github.com/google/cloudprober/web/formatutils"

	surfacerpb "github.com/google/cloudprober/surfacers/proto"
	surfacerspb "github.com/google/cloudprober/surfacers/proto"
)

var (
	userDefinedSurfacers   = make(map[string]Surfacer)
	userDefinedSurfacersMu sync.Mutex
)

// StatusTmpl variable stores the HTML template suitable to generate the
// surfacers' status for cloudprober's /status page. It expects an array of
// SurfacerInfo objects as input.
var StatusTmpl = template.Must(template.New("statusTmpl").Parse(`

    Type
    Name
    Conf

  {{ range . }}

    {{.Type}}
    {{.Name}}

    {{if .Conf}}

{{.Conf}}

    {{else}}
      default
    {{end}}

  {{ end }}

`))

// Default surfacers. These surfacers are enabled if no surfacer is defined.
var defaultSurfacers = []*surfacerpb.SurfacerDef{
	&surfacerpb.SurfacerDef{
		Type: surfacerpb.Type_PROMETHEUS.Enum(),
	},
	&surfacerpb.SurfacerDef{
		Type: surfacerpb.Type_FILE.Enum(),
	},
}

// Surfacer is an interface for all metrics surfacing systems
type Surfacer interface {
	// Function for writing a piece of metric data to a specified metric
	// store (or other location).
	Write(ctx context.Context, em *metrics.EventMetrics)
}

// SurfacerInfo encapsulates a Surfacer and related info.
type SurfacerInfo struct {
	Surfacer
	Type string
	Name string
	Conf string
}

func inferType(s *surfacerpb.SurfacerDef) surfacerspb.Type {
	switch s.Surfacer.(type) {
	case *surfacerpb.SurfacerDef_PrometheusSurfacer:
		return surfacerspb.Type_PROMETHEUS
	case *surfacerpb.SurfacerDef_StackdriverSurfacer:
		return surfacerspb.Type_STACKDRIVER
	case *surfacerpb.SurfacerDef_FileSurfacer:
		return surfacerspb.Type_FILE
	case *surfacerpb.SurfacerDef_PostgresSurfacer:
		return surfacerspb.Type_POSTGRES
	}

	return surfacerspb.Type_NONE
}

// initSurfacer initializes and returns a new surfacer based on the config.
func initSurfacer(s *surfacerpb.SurfacerDef, sType surfacerspb.Type) (Surfacer, interface{}, error) {
	// Create a new logger
	logName := s.GetName()
	if logName == "" {
		logName = strings.ToLower(s.GetType().String())
	}

	// TODO(manugarg): Plumb context here too.
	l, err := logger.New(context.TODO(), logName)
	if err != nil {
		return nil, nil, fmt.Errorf("unable to create cloud logger: %v", err)
	}

	var conf interface{}
	var surfacer Surfacer

	switch sType {
	case surfacerpb.Type_PROMETHEUS:
		surfacer, err = prometheus.New(s.GetPrometheusSurfacer(), l)
		conf = s.GetPrometheusSurfacer()
	case surfacerpb.Type_STACKDRIVER:
		surfacer, err = stackdriver.New(s.GetStackdriverSurfacer(), l)
		conf = s.GetStackdriverSurfacer()
	case surfacerpb.Type_FILE:
		surfacer, err = file.New(s.GetFileSurfacer(), l)
		conf = s.GetFileSurfacer()
	case surfacerpb.Type_POSTGRES:
		surfacer, err = postgres.New(s.GetPostgresSurfacer(), l)
		conf = s.GetPostgresSurfacer()
	case surfacerpb.Type_USER_DEFINED:
		userDefinedSurfacersMu.Lock()
		defer userDefinedSurfacersMu.Unlock()
		surfacer = userDefinedSurfacers[s.GetName()]
		if surfacer == nil {
			return nil, nil, fmt.Errorf("unregistered user defined surfacer: %s", s.GetName())
		}
	default:
		return nil, nil, fmt.Errorf("unknown surfacer type: %s", s.GetType())
	}

	return surfacer, conf, err
}

// Init initializes the surfacers from the config protobufs and returns them as
// a list.
func Init(sDefs []*surfacerpb.SurfacerDef) ([]*SurfacerInfo, error) {
	// If no surfacers are defined, return default surfacers. This behavior
	// can be disabled by explicitly specifying "surfacer {}" in the config.
	if len(sDefs) == 0 {
		sDefs = defaultSurfacers
	}

	var result []*SurfacerInfo
	for _, sDef := range sDefs {
		sType := sDef.GetType()

		if sType == surfacerpb.Type_NONE {
			// Don't do anything if surfacer type is NONE and nothing is defined inside
			// it: for example: "surfacer{}". This is one of the ways to disable
			// surfacers as not adding surfacers at all results in default surfacers
			// being added automatically.
			if sDef.Surfacer == nil {
				continue
			}
			sType = inferType(sDef)
		}

		s, conf, err := initSurfacer(sDef, sType)
		if err != nil {
			return nil, err
		}

		result = append(result, &SurfacerInfo{
			Surfacer: s,
			Type:     sType.String(),
			Name:     sDef.GetName(),
			Conf:     formatutils.ConfToString(conf),
		})
	}
	return result, nil
}

// Register allows you to register a user defined surfacer with cloudprober.
// Example usage:
//	import (
//		"github.com/google/cloudprober"
//		"github.com/google/cloudprober/surfacers"
//	)
//
//	s := &FancySurfacer{}
//	surfacers.Register("fancy_surfacer", s)
//	pr, err := cloudprober.InitFromConfig(*configFile)
//	if err != nil {
//		log.Exitf("Error initializing cloudprober. Err: %v", err)
//	}
func Register(name string, s Surfacer) {
	userDefinedSurfacersMu.Lock()
	defer userDefinedSurfacersMu.Unlock()
	userDefinedSurfacers[name] = s
}
<commit_msg>Fix logger name for surfacers.<commit_after>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/*
Package surfacers is the base package for creating Surfacer objects that are
used for writing metics data to different monitoring services.

Any Surfacer that is created for writing metrics data to a monitor system
should implement the below Surfacer interface and should accept
metrics.EventMetrics object through a Write() call. Each new surfacer should
also plug itself in through the New() method defined here.
*/
package surfacers

import (
	"context"
	"fmt"
	"html/template"
	"strings"
	"sync"

	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/metrics"
	"github.com/google/cloudprober/surfacers/file"
	"github.com/google/cloudprober/surfacers/postgres"
	"github.com/google/cloudprober/surfacers/prometheus"
	"github.com/google/cloudprober/surfacers/stackdriver"
	"github.com/google/cloudprober/web/formatutils"

	surfacerpb "github.com/google/cloudprober/surfacers/proto"
	surfacerspb "github.com/google/cloudprober/surfacers/proto"
)

var (
	userDefinedSurfacers   = make(map[string]Surfacer)
	userDefinedSurfacersMu sync.Mutex
)

// StatusTmpl variable stores the HTML template suitable to generate the
// surfacers' status for cloudprober's /status page. It expects an array of
// SurfacerInfo objects as input.
var StatusTmpl = template.Must(template.New("statusTmpl").Parse(`

    Type
    Name
    Conf

  {{ range . }}

    {{.Type}}
    {{.Name}}

    {{if .Conf}}

{{.Conf}}

    {{else}}
      default
    {{end}}

  {{ end }}

`))

// Default surfacers. These surfacers are enabled if no surfacer is defined.
var defaultSurfacers = []*surfacerpb.SurfacerDef{
	&surfacerpb.SurfacerDef{
		Type: surfacerpb.Type_PROMETHEUS.Enum(),
	},
	&surfacerpb.SurfacerDef{
		Type: surfacerpb.Type_FILE.Enum(),
	},
}

// Surfacer is an interface for all metrics surfacing systems
type Surfacer interface {
	// Function for writing a piece of metric data to a specified metric
	// store (or other location).
	Write(ctx context.Context, em *metrics.EventMetrics)
}

// SurfacerInfo encapsulates a Surfacer and related info.
type SurfacerInfo struct {
	Surfacer
	Type string
	Name string
	Conf string
}

func inferType(s *surfacerpb.SurfacerDef) surfacerspb.Type {
	switch s.Surfacer.(type) {
	case *surfacerpb.SurfacerDef_PrometheusSurfacer:
		return surfacerspb.Type_PROMETHEUS
	case *surfacerpb.SurfacerDef_StackdriverSurfacer:
		return surfacerspb.Type_STACKDRIVER
	case *surfacerpb.SurfacerDef_FileSurfacer:
		return surfacerspb.Type_FILE
	case *surfacerpb.SurfacerDef_PostgresSurfacer:
		return surfacerspb.Type_POSTGRES
	}

	return surfacerspb.Type_NONE
}

// initSurfacer initializes and returns a new surfacer based on the config.
func initSurfacer(s *surfacerpb.SurfacerDef, sType surfacerspb.Type) (Surfacer, interface{}, error) {
	// Create a new logger
	logName := s.GetName()
	if logName == "" {
		logName = strings.ToLower(s.GetType().String())
	}

	l, err := logger.NewCloudproberLog(logName)
	if err != nil {
		return nil, nil, fmt.Errorf("unable to create cloud logger: %v", err)
	}

	var conf interface{}
	var surfacer Surfacer

	switch sType {
	case surfacerpb.Type_PROMETHEUS:
		surfacer, err = prometheus.New(s.GetPrometheusSurfacer(), l)
		conf = s.GetPrometheusSurfacer()
	case surfacerpb.Type_STACKDRIVER:
		surfacer, err = stackdriver.New(s.GetStackdriverSurfacer(), l)
		conf = s.GetStackdriverSurfacer()
	case surfacerpb.Type_FILE:
		surfacer, err = file.New(s.GetFileSurfacer(), l)
		conf = s.GetFileSurfacer()
	case surfacerpb.Type_POSTGRES:
		surfacer, err = postgres.New(s.GetPostgresSurfacer(), l)
		conf = s.GetPostgresSurfacer()
	case surfacerpb.Type_USER_DEFINED:
		userDefinedSurfacersMu.Lock()
		defer userDefinedSurfacersMu.Unlock()
		surfacer = userDefinedSurfacers[s.GetName()]
		if surfacer == nil {
			return nil, nil, fmt.Errorf("unregistered user defined surfacer: %s", s.GetName())
		}
	default:
		return nil, nil, fmt.Errorf("unknown surfacer type: %s", s.GetType())
	}

	return surfacer, conf, err
}

// Init initializes the surfacers from the config protobufs and returns them as
// a list.
func Init(sDefs []*surfacerpb.SurfacerDef) ([]*SurfacerInfo, error) {
	// If no surfacers are defined, return default surfacers. This behavior
	// can be disabled by explicitly specifying "surfacer {}" in the config.
	if len(sDefs) == 0 {
		sDefs = defaultSurfacers
	}

	var result []*SurfacerInfo
	for _, sDef := range sDefs {
		sType := sDef.GetType()

		if sType == surfacerpb.Type_NONE {
			// Don't do anything if surfacer type is NONE and nothing is defined inside
			// it: for example: "surfacer{}". This is one of the ways to disable
			// surfacers as not adding surfacers at all results in default surfacers
			// being added automatically.
			if sDef.Surfacer == nil {
				continue
			}
			sType = inferType(sDef)
		}

		s, conf, err := initSurfacer(sDef, sType)
		if err != nil {
			return nil, err
		}

		result = append(result, &SurfacerInfo{
			Surfacer: s,
			Type:     sType.String(),
			Name:     sDef.GetName(),
			Conf:     formatutils.ConfToString(conf),
		})
	}
	return result, nil
}

// Register allows you to register a user defined surfacer with cloudprober.
// Example usage:
//	import (
//		"github.com/google/cloudprober"
//		"github.com/google/cloudprober/surfacers"
//	)
//
//	s := &FancySurfacer{}
//	surfacers.Register("fancy_surfacer", s)
//	pr, err := cloudprober.InitFromConfig(*configFile)
//	if err != nil {
//		log.Exitf("Error initializing cloudprober. Err: %v", err)
//	}
func Register(name string, s Surfacer) {
	userDefinedSurfacersMu.Lock()
	defer userDefinedSurfacersMu.Unlock()
	userDefinedSurfacers[name] = s
}
<|endoftext|>


<commit_before>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2016

package instana

import (
	"errors"
	"net/http"
	"strings"

	"github.com/instana/go-sensor/w3ctrace"
	ot "github.com/opentracing/opentracing-go"
)

// Instana header constants
const (
	// FieldT Trace ID header
	FieldT = "x-instana-t"
	// FieldS Span ID header
	FieldS = "x-instana-s"
	// FieldL Level header
	FieldL = "x-instana-l"
	// FieldB OT Baggage header
	FieldB = "x-instana-b-"
	// FieldSynthetic if set to 1, marks the call as synthetic, e.g.
	// a healthcheck request
	FieldSynthetic = "x-instana-synthetic"
)

func injectTraceContext(sc SpanContext, opaqueCarrier interface{}) error {
	roCarrier, ok := opaqueCarrier.(ot.TextMapReader)
	if !ok {
		return ot.ErrInvalidCarrier
	}

	// Handle pre-existing case-sensitive keys
	exstfieldT := FieldT
	exstfieldS := FieldS
	exstfieldL := FieldL
	exstfieldB := FieldB

	roCarrier.ForeachKey(func(k, v string) error {
		switch strings.ToLower(k) {
		case FieldT:
			exstfieldT = k
		case FieldS:
			exstfieldS = k
		case FieldL:
			exstfieldL = k
		default:
			if strings.HasPrefix(strings.ToLower(k), FieldB) {
				exstfieldB = string([]rune(k)[:len(FieldB)])
			}
		}
		return nil
	})

	carrier, ok := opaqueCarrier.(ot.TextMapWriter)
	if !ok {
		return ot.ErrInvalidCarrier
	}

	if c, ok := opaqueCarrier.(ot.HTTPHeadersCarrier); ok {
		// Even though the godoc claims that the key passed to (*http.Header).Set()
		// is case-insensitive, it actually normalizes it using textproto.CanonicalMIMEHeaderKey()
		// before populating the value. As a result headers with non-canonical will not be
		// overwritten with a new value. This is only the case if header names were set while
		// initializing the http.Header instance, i.e.
		//     h := http.Headers{"X-InStAnA-T": {"abc123"}}
		// and does not apply to a common case when requests are being created using http.NewRequest()
		// or http.ReadRequest() that call (*http.Header).Set() to set header values.
		h := http.Header(c)
		delete(h, exstfieldT)
		delete(h, exstfieldS)
		delete(h, exstfieldL)

		for key := range h {
			if strings.HasPrefix(strings.ToLower(key), FieldB) {
				delete(h, key)
			}
		}

		addW3CTraceContext(h, sc)
		addEUMHeaders(h, sc)
	}

	if !sc.Suppressed {
		carrier.Set(exstfieldT, FormatID(sc.TraceID))
		carrier.Set(exstfieldS, FormatID(sc.SpanID))
	} else {
		// remove trace context keys from the carrier
		switch c := opaqueCarrier.(type) {
		case ot.HTTPHeadersCarrier:
			h := http.Header(c)
			h.Del(exstfieldT)
			h.Del(exstfieldS)
		case ot.TextMapCarrier:
			delete(c, exstfieldT)
			delete(c, exstfieldS)
		case interface{ RemoveAll() }:
			// in case carrier has the RemoveAll() method that wipes all trace
			// headers, for example the instasarama.ProducerMessagCarrier, we
			// use it to remove the context of a suppressed trace
			c.RemoveAll()
		}
	}

	carrier.Set(exstfieldL, formatLevel(sc))

	for k, v := range sc.Baggage {
		carrier.Set(exstfieldB+k, v)
	}

	return nil
}

// This method searches for Instana headers (FieldT, FieldS, FieldL and header with name prefixed with FieldB)
// and try to parse their values. It also tries to extract w3 context and assign it inside returned object. W3 context
// will be propagated further and can be used as a fallback.
func extractTraceContext(opaqueCarrier interface{}) (SpanContext, error) {
	spanContext := SpanContext{
		Baggage: make(map[string]string),
	}

	carrier, ok := opaqueCarrier.(ot.TextMapReader)
	if !ok {
		return spanContext, ot.ErrInvalidCarrier
	}

	if c, ok := opaqueCarrier.(ot.HTTPHeadersCarrier); ok {
		pickupW3CTraceContext(http.Header(c), &spanContext)
	}

	// Iterate over the headers, look for Instana headers and try to parse them.
	// In case of error interrupt, iteration and return it.
	err := carrier.ForeachKey(func(k, v string) error {
		var err error

		switch strings.ToLower(k) {
		case FieldT:
			spanContext.TraceIDHi, spanContext.TraceID, err = ParseLongID(v)
			if err != nil {
				sensor.logger.Debug("extract trace context ", FieldT, "=", v, " : ", err)
				return ot.ErrSpanContextCorrupted
			}
		case FieldS:
			spanContext.SpanID, err = ParseID(v)
			if err != nil {
				sensor.logger.Debug("extract trace context ", FieldT, "=", v, " : ", err)
				return ot.ErrSpanContextCorrupted
			}
		case FieldL:
			// When FieldL is present and equals to "0", then spanContext is suppressed.
			// In addition to that non-empty correlation data may be extracted.
			suppressed, corrData, err := parseLevel(v)
			if err != nil {
				sensor.logger.Info("failed to parse ", k, ": ", err, " (", v, ")")
				// use defaults
				suppressed, corrData = false, EUMCorrelationData{}
			}

			spanContext.Suppressed = suppressed
			if !spanContext.Suppressed {
				spanContext.Correlation = corrData
			}
		default:
			if strings.HasPrefix(strings.ToLower(k), FieldB) {
				// preserve original case of the baggage key
				spanContext.Baggage[k[len(FieldB):]] = v
			}
		}

		return nil
	})
	if err != nil {
		return spanContext, err
	}

	// reset the trace IDs if a correlation ID has been provided
	if spanContext.Correlation.ID != "" {
		spanContext.TraceIDHi, spanContext.TraceID, spanContext.SpanID = 0, 0, 0

		return spanContext, nil
	}

	if spanContext.IsZero() {
		return spanContext, ot.ErrSpanContextNotFound
	}

	// When the context is not suppressed and one of Instana ID headers set.
	if !spanContext.Suppressed &&
		(spanContext.SpanID == 0 != (spanContext.TraceIDHi == 0 && spanContext.TraceID == 0)) {
		sensor.logger.Debug("broken Instana trace context:",
			" SpanID=", FormatID(spanContext.SpanID),
			" TraceID=", FormatLongID(spanContext.TraceIDHi, spanContext.TraceID))

		// Check if w3 context was found or not
		if spanContext.W3CContext.IsZero() {
			return spanContext, ot.ErrSpanContextCorrupted
		} else {
			// Return SpanContext with w3 context, ignore other values
			return SpanContext{
				W3CContext: spanContext.W3CContext,
			}, nil
		}

	}

	return spanContext, nil
}

func addW3CTraceContext(h http.Header, sc SpanContext) {
	traceID, spanID := FormatLongID(sc.TraceIDHi, sc.TraceID), FormatID(sc.SpanID)
	trCtx := sc.W3CContext

	// check for an existing w3c trace
	if trCtx.IsZero() {
		// initiate trace if none
		trCtx = w3ctrace.New(w3ctrace.Parent{
			Version:  w3ctrace.Version_Max,
			TraceID:  traceID,
			ParentID: spanID,
			Flags: w3ctrace.Flags{
				Sampled: !sc.Suppressed,
			},
		})
	}

	// update the traceparent parent ID
	p := trCtx.Parent()
	p.ParentID = spanID
	// sync the traceparent `sampled` flags with the X-Instana-L value
	p.Flags.Sampled = !sc.Suppressed

	trCtx.RawParent = p.String()

	// participate in w3c trace context if tracing is enabled
	if !sc.Suppressed {
		// propagate truncated trace ID downstream
		trCtx.RawState = trCtx.State().Add(w3ctrace.VendorInstana, FormatID(sc.TraceID)+";"+spanID).String()
	}

	w3ctrace.Inject(trCtx, h)
}

func pickupW3CTraceContext(h http.Header, sc *SpanContext) {
	trCtx, err := w3ctrace.Extract(h)
	if err != nil {
		return
	}
	sc.W3CContext = trCtx
}

func addEUMHeaders(h http.Header, sc SpanContext) {
	// Preserve original Server-Timing header values by combining them into a comma-separated list
	st := append(h["Server-Timing"], "intid;desc="+FormatID(sc.TraceID))
	h.Set("Server-Timing", strings.Join(st, ", "))
}

var errMalformedHeader = errors.New("malformed header value")

func parseLevel(s string) (bool, EUMCorrelationData, error) {
	const (
		levelState uint8 = iota
		partSeparatorState
		correlationPartState
		correlationTypeState
		correlationIDState
		finalState
	)

	if s == "" {
		return false, EUMCorrelationData{}, nil
	}

	var (
		typeInd                 int
		state                   uint8
		level, corrType, corrID string
	)
PARSE:
	for ptr := 0; state != finalState && ptr < len(s); ptr++ {
		switch state {
		case levelState: // looking for 0 or 1
			level = s[ptr : ptr+1]

			if level != "0" && level != "1" {
				break PARSE
			}

			if ptr == len(s)-1 { // no correlation ID provided
				state = finalState
			} else {
				state = partSeparatorState
			}
		case partSeparatorState: // skip OWS while looking for ','
			switch s[ptr] {
			case ' ', '\t': // advance
			case ',':
				state = correlationPartState
			default:
				break PARSE
			}
		case correlationPartState: // skip OWS while searching for 'correlationType=' prefix
			switch {
			case s[ptr] == ' ' || s[ptr] == '\t': // advance
			case strings.HasPrefix(s[ptr:], "correlationType="):
				ptr += 15 // advance to the end of prefix
				typeInd = ptr + 1
				state = correlationTypeState
			default:
				break PARSE
			}
		case correlationTypeState: // skip OWS while looking for ';'
			switch s[ptr] {
			case ' ', '\t': // possibly trailing OWS, advance
			case ';':
				state = correlationIDState
			default:
				corrType = s[typeInd : ptr+1]
			}
		case correlationIDState: //  skip OWS while searching for 'correlationId=' prefix
			switch {
			case s[ptr] == ' ' || s[ptr] == '\t': // leading OWS, advance
			case strings.HasPrefix(s[ptr:], "correlationId="):
				ptr += 14
				corrID = s[ptr:]
				state = finalState
			default:
				break PARSE
			}
		default:
			break PARSE
		}
	}

	if state != finalState {
		return false, EUMCorrelationData{}, errMalformedHeader
	}

	return level == "0", EUMCorrelationData{Type: corrType, ID: corrID}, nil
}

func formatLevel(sc SpanContext) string {
	if sc.Suppressed {
		return "0"
	}

	return "1"
}
<commit_msg>Update propagation.go<commit_after>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2016

package instana

import (
	"errors"
	"net/http"
	"strings"

	"github.com/instana/go-sensor/w3ctrace"
	ot "github.com/opentracing/opentracing-go"
)

// Instana header constants
const (
	// FieldT Trace ID header
	FieldT = "x-instana-t"
	// FieldS Span ID header
	FieldS = "x-instana-s"
	// FieldL Level header
	FieldL = "x-instana-l"
	// FieldB OT Baggage header
	FieldB = "x-instana-b-"
	// FieldSynthetic if set to 1, marks the call as synthetic, e.g.
	// a healthcheck request
	FieldSynthetic = "x-instana-synthetic"
)

func injectTraceContext(sc SpanContext, opaqueCarrier interface{}) error {
	roCarrier, ok := opaqueCarrier.(ot.TextMapReader)
	if !ok {
		return ot.ErrInvalidCarrier
	}

	// Handle pre-existing case-sensitive keys
	exstfieldT := FieldT
	exstfieldS := FieldS
	exstfieldL := FieldL
	exstfieldB := FieldB

	roCarrier.ForeachKey(func(k, v string) error {
		switch strings.ToLower(k) {
		case FieldT:
			exstfieldT = k
		case FieldS:
			exstfieldS = k
		case FieldL:
			exstfieldL = k
		default:
			if strings.HasPrefix(strings.ToLower(k), FieldB) {
				exstfieldB = string([]rune(k)[:len(FieldB)])
			}
		}
		return nil
	})

	carrier, ok := opaqueCarrier.(ot.TextMapWriter)
	if !ok {
		return ot.ErrInvalidCarrier
	}

	if c, ok := opaqueCarrier.(ot.HTTPHeadersCarrier); ok {
		// Even though the godoc claims that the key passed to (*http.Header).Set()
		// is case-insensitive, it actually normalizes it using textproto.CanonicalMIMEHeaderKey()
		// before populating the value. As a result headers with non-canonical will not be
		// overwritten with a new value. This is only the case if header names were set while
		// initializing the http.Header instance, i.e.
		//     h := http.Headers{"X-InStAnA-T": {"abc123"}}
		// and does not apply to a common case when requests are being created using http.NewRequest()
		// or http.ReadRequest() that call (*http.Header).Set() to set header values.
		h := http.Header(c)
		delete(h, exstfieldT)
		delete(h, exstfieldS)
		delete(h, exstfieldL)

		for key := range h {
			if strings.HasPrefix(strings.ToLower(key), FieldB) {
				delete(h, key)
			}
		}

		addW3CTraceContext(h, sc)
		addEUMHeaders(h, sc)
	}

	if !sc.Suppressed {
		carrier.Set(exstfieldT, FormatID(sc.TraceID))
		carrier.Set(exstfieldS, FormatID(sc.SpanID))
	} else {
		// remove trace context keys from the carrier
		switch c := opaqueCarrier.(type) {
		case ot.HTTPHeadersCarrier:
			h := http.Header(c)
			h.Del(exstfieldT)
			h.Del(exstfieldS)
		case ot.TextMapCarrier:
			delete(c, exstfieldT)
			delete(c, exstfieldS)
		case interface{ RemoveAll() }:
			// in case carrier has the RemoveAll() method that wipes all trace
			// headers, for example the instasarama.ProducerMessagCarrier, we
			// use it to remove the context of a suppressed trace
			c.RemoveAll()
		}
	}

	carrier.Set(exstfieldL, formatLevel(sc))

	for k, v := range sc.Baggage {
		carrier.Set(exstfieldB+k, v)
	}

	return nil
}

// This method searches for Instana headers (FieldT, FieldS, FieldL and header with name prefixed with FieldB)
// and try to parse their values. It also tries to extract W3C context and assign it inside returned object. W3C context
// will be propagated further and can be used as a fallback.
func extractTraceContext(opaqueCarrier interface{}) (SpanContext, error) {
	spanContext := SpanContext{
		Baggage: make(map[string]string),
	}

	carrier, ok := opaqueCarrier.(ot.TextMapReader)
	if !ok {
		return spanContext, ot.ErrInvalidCarrier
	}

	if c, ok := opaqueCarrier.(ot.HTTPHeadersCarrier); ok {
		pickupW3CTraceContext(http.Header(c), &spanContext)
	}

	// Iterate over the headers, look for Instana headers and try to parse them.
	// In case of error interrupt, iteration and return it.
	err := carrier.ForeachKey(func(k, v string) error {
		var err error

		switch strings.ToLower(k) {
		case FieldT:
			spanContext.TraceIDHi, spanContext.TraceID, err = ParseLongID(v)
			if err != nil {
				sensor.logger.Debug("extract trace context ", FieldT, "=", v, " : ", err)
				return ot.ErrSpanContextCorrupted
			}
		case FieldS:
			spanContext.SpanID, err = ParseID(v)
			if err != nil {
				sensor.logger.Debug("extract trace context ", FieldT, "=", v, " : ", err)
				return ot.ErrSpanContextCorrupted
			}
		case FieldL:
			// When FieldL is present and equals to "0", then spanContext is suppressed.
			// In addition to that non-empty correlation data may be extracted.
			suppressed, corrData, err := parseLevel(v)
			if err != nil {
				sensor.logger.Info("failed to parse ", k, ": ", err, " (", v, ")")
				// use defaults
				suppressed, corrData = false, EUMCorrelationData{}
			}

			spanContext.Suppressed = suppressed
			if !spanContext.Suppressed {
				spanContext.Correlation = corrData
			}
		default:
			if strings.HasPrefix(strings.ToLower(k), FieldB) {
				// preserve original case of the baggage key
				spanContext.Baggage[k[len(FieldB):]] = v
			}
		}

		return nil
	})
	if err != nil {
		return spanContext, err
	}

	// reset the trace IDs if a correlation ID has been provided
	if spanContext.Correlation.ID != "" {
		spanContext.TraceIDHi, spanContext.TraceID, spanContext.SpanID = 0, 0, 0

		return spanContext, nil
	}

	if spanContext.IsZero() {
		return spanContext, ot.ErrSpanContextNotFound
	}

	// When the context is not suppressed and one of Instana ID headers set.
	if !spanContext.Suppressed &&
		(spanContext.SpanID == 0 != (spanContext.TraceIDHi == 0 && spanContext.TraceID == 0)) {
		sensor.logger.Debug("broken Instana trace context:",
			" SpanID=", FormatID(spanContext.SpanID),
			" TraceID=", FormatLongID(spanContext.TraceIDHi, spanContext.TraceID))

		// Check if w3 context was found or not
		if spanContext.W3CContext.IsZero() {
			return spanContext, ot.ErrSpanContextCorrupted
		} else {
			// Return SpanContext with w3 context, ignore other values
			return SpanContext{
				W3CContext: spanContext.W3CContext,
			}, nil
		}

	}

	return spanContext, nil
}

func addW3CTraceContext(h http.Header, sc SpanContext) {
	traceID, spanID := FormatLongID(sc.TraceIDHi, sc.TraceID), FormatID(sc.SpanID)
	trCtx := sc.W3CContext

	// check for an existing w3c trace
	if trCtx.IsZero() {
		// initiate trace if none
		trCtx = w3ctrace.New(w3ctrace.Parent{
			Version:  w3ctrace.Version_Max,
			TraceID:  traceID,
			ParentID: spanID,
			Flags: w3ctrace.Flags{
				Sampled: !sc.Suppressed,
			},
		})
	}

	// update the traceparent parent ID
	p := trCtx.Parent()
	p.ParentID = spanID
	// sync the traceparent `sampled` flags with the X-Instana-L value
	p.Flags.Sampled = !sc.Suppressed

	trCtx.RawParent = p.String()

	// participate in w3c trace context if tracing is enabled
	if !sc.Suppressed {
		// propagate truncated trace ID downstream
		trCtx.RawState = trCtx.State().Add(w3ctrace.VendorInstana, FormatID(sc.TraceID)+";"+spanID).String()
	}

	w3ctrace.Inject(trCtx, h)
}

func pickupW3CTraceContext(h http.Header, sc *SpanContext) {
	trCtx, err := w3ctrace.Extract(h)
	if err != nil {
		return
	}
	sc.W3CContext = trCtx
}

func addEUMHeaders(h http.Header, sc SpanContext) {
	// Preserve original Server-Timing header values by combining them into a comma-separated list
	st := append(h["Server-Timing"], "intid;desc="+FormatID(sc.TraceID))
	h.Set("Server-Timing", strings.Join(st, ", "))
}

var errMalformedHeader = errors.New("malformed header value")

func parseLevel(s string) (bool, EUMCorrelationData, error) {
	const (
		levelState uint8 = iota
		partSeparatorState
		correlationPartState
		correlationTypeState
		correlationIDState
		finalState
	)

	if s == "" {
		return false, EUMCorrelationData{}, nil
	}

	var (
		typeInd                 int
		state                   uint8
		level, corrType, corrID string
	)
PARSE:
	for ptr := 0; state != finalState && ptr < len(s); ptr++ {
		switch state {
		case levelState: // looking for 0 or 1
			level = s[ptr : ptr+1]

			if level != "0" && level != "1" {
				break PARSE
			}

			if ptr == len(s)-1 { // no correlation ID provided
				state = finalState
			} else {
				state = partSeparatorState
			}
		case partSeparatorState: // skip OWS while looking for ','
			switch s[ptr] {
			case ' ', '\t': // advance
			case ',':
				state = correlationPartState
			default:
				break PARSE
			}
		case correlationPartState: // skip OWS while searching for 'correlationType=' prefix
			switch {
			case s[ptr] == ' ' || s[ptr] == '\t': // advance
			case strings.HasPrefix(s[ptr:], "correlationType="):
				ptr += 15 // advance to the end of prefix
				typeInd = ptr + 1
				state = correlationTypeState
			default:
				break PARSE
			}
		case correlationTypeState: // skip OWS while looking for ';'
			switch s[ptr] {
			case ' ', '\t': // possibly trailing OWS, advance
			case ';':
				state = correlationIDState
			default:
				corrType = s[typeInd : ptr+1]
			}
		case correlationIDState: //  skip OWS while searching for 'correlationId=' prefix
			switch {
			case s[ptr] == ' ' || s[ptr] == '\t': // leading OWS, advance
			case strings.HasPrefix(s[ptr:], "correlationId="):
				ptr += 14
				corrID = s[ptr:]
				state = finalState
			default:
				break PARSE
			}
		default:
			break PARSE
		}
	}

	if state != finalState {
		return false, EUMCorrelationData{}, errMalformedHeader
	}

	return level == "0", EUMCorrelationData{Type: corrType, ID: corrID}, nil
}

func formatLevel(sc SpanContext) string {
	if sc.Suppressed {
		return "0"
	}

	return "1"
}
<|endoftext|>


<commit_before>// Package provider holds the different provider implementation.
package provider

import (
	"errors"
	"fmt"
	"strings"
	"text/template"
	"time"

	"github.com/BurntSushi/ty/fun"
	"github.com/cenk/backoff"
	"github.com/containous/traefik/job"
	"github.com/containous/traefik/log"
	"github.com/containous/traefik/safe"
	"github.com/containous/traefik/types"
	"github.com/docker/libkv"
	"github.com/docker/libkv/store"
)

// Kv holds common configurations of key-value providers.
type Kv struct {
	BaseProvider `mapstructure:",squash"`
	Endpoint     string     `description:"Comma sepparated server endpoints"`
	Prefix       string     `description:"Prefix used for KV store"`
	TLS          *ClientTLS `description:"Enable TLS support"`
	storeType    store.Backend
	kvclient     store.Store
}

func (provider *Kv) createStore() (store.Store, error) {
	storeConfig := &store.Config{
		ConnectionTimeout: 30 * time.Second,
		Bucket:            "traefik",
	}

	if provider.TLS != nil {
		var err error
		storeConfig.TLS, err = provider.TLS.CreateTLSConfig()
		if err != nil {
			return nil, err
		}
	}
	return libkv.NewStore(
		provider.storeType,
		strings.Split(provider.Endpoint, ","),
		storeConfig,
	)
}

func (provider *Kv) watchKv(configurationChan chan<- types.ConfigMessage, prefix string, stop chan bool) error {
	operation := func() error {
		events, err := provider.kvclient.WatchTree(provider.Prefix, make(chan struct{}))
		if err != nil {
			return fmt.Errorf("Failed to KV WatchTree: %v", err)
		}
		for {
			select {
			case <-stop:
				return nil
			case _, ok := <-events:
				if !ok {
					return errors.New("watchtree channel closed")
				}
				configuration := provider.loadConfig()
				if configuration != nil {
					configurationChan <- types.ConfigMessage{
						ProviderName:  string(provider.storeType),
						Configuration: configuration,
					}
				}
			}
		}
	}

	notify := func(err error, time time.Duration) {
		log.Errorf("KV connection error: %+v, retrying in %s", err, time)
	}
	err := backoff.RetryNotify(safe.OperationWithRecover(operation), job.NewBackOff(backoff.NewExponentialBackOff()), notify)
	if err != nil {
		return fmt.Errorf("Cannot connect to KV server: %v", err)
	}
	return nil
}

func (provider *Kv) provide(configurationChan chan<- types.ConfigMessage, pool *safe.Pool, constraints types.Constraints) error {
	provider.Constraints = append(provider.Constraints, constraints...)
	operation := func() error {
		if _, err := provider.kvclient.Exists("qmslkjdfmqlskdjfmqlksjazçueznbvbwzlkajzebvkwjdcqmlsfj"); err != nil {
			return fmt.Errorf("Failed to test KV store connection: %v", err)
		}
		if provider.Watch {
			pool.Go(func(stop chan bool) {
				err := provider.watchKv(configurationChan, provider.Prefix, stop)
				if err != nil {
					log.Errorf("Cannot watch KV store: %v", err)
				}
			})
		}
		configuration := provider.loadConfig()
		configurationChan <- types.ConfigMessage{
			ProviderName:  string(provider.storeType),
			Configuration: configuration,
		}
		return nil
	}
	notify := func(err error, time time.Duration) {
		log.Errorf("KV connection error: %+v, retrying in %s", err, time)
	}
	err := backoff.RetryNotify(safe.OperationWithRecover(operation), job.NewBackOff(backoff.NewExponentialBackOff()), notify)
	if err != nil {
		return fmt.Errorf("Cannot connect to KV server: %v", err)
	}
	return nil
}

func (provider *Kv) loadConfig() *types.Configuration {
	templateObjects := struct {
		Prefix string
	}{
		// Allow `/traefik/alias` to superesede `provider.Prefix`
		strings.TrimSuffix(provider.get(provider.Prefix, provider.Prefix+"/alias"), "/"),
	}

	var KvFuncMap = template.FuncMap{
		"List":        provider.list,
		"ListServers": provider.listServers,
		"Get":         provider.get,
		"SplitGet":    provider.splitGet,
		"Last":        provider.last,
	}

	configuration, err := provider.getConfiguration("templates/kv.tmpl", KvFuncMap, templateObjects)
	if err != nil {
		log.Error(err)
	}

	for key, frontend := range configuration.Frontends {
		if _, ok := configuration.Backends[frontend.Backend]; ok == false {
			delete(configuration.Frontends, key)
		}
	}

	return configuration
}

func (provider *Kv) list(keys ...string) []string {
	joinedKeys := strings.Join(keys, "")
	keysPairs, err := provider.kvclient.List(joinedKeys)
	if err != nil {
		log.Debugf("Cannot get keys %s %s ", joinedKeys, err)
		return nil
	}
	directoryKeys := make(map[string]string)
	for _, key := range keysPairs {
		directory := strings.Split(strings.TrimPrefix(key.Key, joinedKeys), "/")[0]
		directoryKeys[directory] = joinedKeys + directory
	}
	return fun.Values(directoryKeys).([]string)
}

func (provider *Kv) listServers(backend string) []string {
	serverNames := provider.list(backend, "/servers/")
	return fun.Filter(func(serverName string) bool {
		if _, err := provider.kvclient.Get(fmt.Sprint(serverName, "/url")); err != nil {
			return false
		}
		return provider.checkConstraints(serverName, "/tags")
	}, serverNames).([]string)
}

func (provider *Kv) get(defaultValue string, keys ...string) string {
	joinedKeys := strings.Join(keys, "")
	keyPair, err := provider.kvclient.Get(strings.TrimPrefix(joinedKeys, "/"))
	if err != nil {
		log.Debugf("Cannot get key %s %s, setting default %s", joinedKeys, err, defaultValue)
		return defaultValue
	} else if keyPair == nil {
		log.Debugf("Cannot get key %s, setting default %s", joinedKeys, defaultValue)
		return defaultValue
	}
	return string(keyPair.Value)
}

func (provider *Kv) splitGet(keys ...string) []string {
	joinedKeys := strings.Join(keys, "")
	keyPair, err := provider.kvclient.Get(joinedKeys)
	if err != nil {
		log.Debugf("Cannot get key %s %s, setting default empty", joinedKeys, err)
		return []string{}
	} else if keyPair == nil {
		log.Debugf("Cannot get key %s, setting default %empty", joinedKeys)
		return []string{}
	}
	return strings.Split(string(keyPair.Value), ",")
}

func (provider *Kv) last(key string) string {
	splittedKey := strings.Split(key, "/")
	return splittedKey[len(splittedKey)-1]
}

func (provider *Kv) checkConstraints(keys ...string) bool {
	joinedKeys := strings.Join(keys, "")
	keyPair, err := provider.kvclient.Get(joinedKeys)

	value := ""
	if err == nil && keyPair != nil && keyPair.Value != nil {
		value = string(keyPair.Value)
	}

	constraintTags := strings.Split(value, ",")
	ok, failingConstraint := provider.MatchConstraints(constraintTags)
	if ok == false {
		if failingConstraint != nil {
			log.Debugf("Constraint %v not matching with following tags: %v", failingConstraint.String(), value)
		}
		return false
	}
	return true
}
<commit_msg>kv: Log error when checking existence of server url key<commit_after>// Package provider holds the different provider implementation.
package provider

import (
	"errors"
	"fmt"
	"strings"
	"text/template"
	"time"

	"github.com/BurntSushi/ty/fun"
	"github.com/cenk/backoff"
	"github.com/containous/traefik/job"
	"github.com/containous/traefik/log"
	"github.com/containous/traefik/safe"
	"github.com/containous/traefik/types"
	"github.com/docker/libkv"
	"github.com/docker/libkv/store"
)

// Kv holds common configurations of key-value providers.
type Kv struct {
	BaseProvider `mapstructure:",squash"`
	Endpoint     string     `description:"Comma sepparated server endpoints"`
	Prefix       string     `description:"Prefix used for KV store"`
	TLS          *ClientTLS `description:"Enable TLS support"`
	storeType    store.Backend
	kvclient     store.Store
}

func (provider *Kv) createStore() (store.Store, error) {
	storeConfig := &store.Config{
		ConnectionTimeout: 30 * time.Second,
		Bucket:            "traefik",
	}

	if provider.TLS != nil {
		var err error
		storeConfig.TLS, err = provider.TLS.CreateTLSConfig()
		if err != nil {
			return nil, err
		}
	}
	return libkv.NewStore(
		provider.storeType,
		strings.Split(provider.Endpoint, ","),
		storeConfig,
	)
}

func (provider *Kv) watchKv(configurationChan chan<- types.ConfigMessage, prefix string, stop chan bool) error {
	operation := func() error {
		events, err := provider.kvclient.WatchTree(provider.Prefix, make(chan struct{}))
		if err != nil {
			return fmt.Errorf("Failed to KV WatchTree: %v", err)
		}
		for {
			select {
			case <-stop:
				return nil
			case _, ok := <-events:
				if !ok {
					return errors.New("watchtree channel closed")
				}
				configuration := provider.loadConfig()
				if configuration != nil {
					configurationChan <- types.ConfigMessage{
						ProviderName:  string(provider.storeType),
						Configuration: configuration,
					}
				}
			}
		}
	}

	notify := func(err error, time time.Duration) {
		log.Errorf("KV connection error: %+v, retrying in %s", err, time)
	}
	err := backoff.RetryNotify(safe.OperationWithRecover(operation), job.NewBackOff(backoff.NewExponentialBackOff()), notify)
	if err != nil {
		return fmt.Errorf("Cannot connect to KV server: %v", err)
	}
	return nil
}

func (provider *Kv) provide(configurationChan chan<- types.ConfigMessage, pool *safe.Pool, constraints types.Constraints) error {
	provider.Constraints = append(provider.Constraints, constraints...)
	operation := func() error {
		if _, err := provider.kvclient.Exists("qmslkjdfmqlskdjfmqlksjazçueznbvbwzlkajzebvkwjdcqmlsfj"); err != nil {
			return fmt.Errorf("Failed to test KV store connection: %v", err)
		}
		if provider.Watch {
			pool.Go(func(stop chan bool) {
				err := provider.watchKv(configurationChan, provider.Prefix, stop)
				if err != nil {
					log.Errorf("Cannot watch KV store: %v", err)
				}
			})
		}
		configuration := provider.loadConfig()
		configurationChan <- types.ConfigMessage{
			ProviderName:  string(provider.storeType),
			Configuration: configuration,
		}
		return nil
	}
	notify := func(err error, time time.Duration) {
		log.Errorf("KV connection error: %+v, retrying in %s", err, time)
	}
	err := backoff.RetryNotify(safe.OperationWithRecover(operation), job.NewBackOff(backoff.NewExponentialBackOff()), notify)
	if err != nil {
		return fmt.Errorf("Cannot connect to KV server: %v", err)
	}
	return nil
}

func (provider *Kv) loadConfig() *types.Configuration {
	templateObjects := struct {
		Prefix string
	}{
		// Allow `/traefik/alias` to superesede `provider.Prefix`
		strings.TrimSuffix(provider.get(provider.Prefix, provider.Prefix+"/alias"), "/"),
	}

	var KvFuncMap = template.FuncMap{
		"List":        provider.list,
		"ListServers": provider.listServers,
		"Get":         provider.get,
		"SplitGet":    provider.splitGet,
		"Last":        provider.last,
	}

	configuration, err := provider.getConfiguration("templates/kv.tmpl", KvFuncMap, templateObjects)
	if err != nil {
		log.Error(err)
	}

	for key, frontend := range configuration.Frontends {
		if _, ok := configuration.Backends[frontend.Backend]; ok == false {
			delete(configuration.Frontends, key)
		}
	}

	return configuration
}

func (provider *Kv) list(keys ...string) []string {
	joinedKeys := strings.Join(keys, "")
	keysPairs, err := provider.kvclient.List(joinedKeys)
	if err != nil {
		log.Debugf("Cannot get keys %s %s ", joinedKeys, err)
		return nil
	}
	directoryKeys := make(map[string]string)
	for _, key := range keysPairs {
		directory := strings.Split(strings.TrimPrefix(key.Key, joinedKeys), "/")[0]
		directoryKeys[directory] = joinedKeys + directory
	}
	return fun.Values(directoryKeys).([]string)
}

func (provider *Kv) listServers(backend string) []string {
	serverNames := provider.list(backend, "/servers/")
	return fun.Filter(func(serverName string) bool {
		key := fmt.Sprint(serverName, "/url")
		if _, err := provider.kvclient.Get(key); err != nil {
			if err != store.ErrKeyNotFound {
				log.Errorf("Failed to retrieve value for key %s: %s", key, err)
			}
			return false
		}
		return provider.checkConstraints(serverName, "/tags")
	}, serverNames).([]string)
}

func (provider *Kv) get(defaultValue string, keys ...string) string {
	joinedKeys := strings.Join(keys, "")
	keyPair, err := provider.kvclient.Get(strings.TrimPrefix(joinedKeys, "/"))
	if err != nil {
		log.Debugf("Cannot get key %s %s, setting default %s", joinedKeys, err, defaultValue)
		return defaultValue
	} else if keyPair == nil {
		log.Debugf("Cannot get key %s, setting default %s", joinedKeys, defaultValue)
		return defaultValue
	}
	return string(keyPair.Value)
}

func (provider *Kv) splitGet(keys ...string) []string {
	joinedKeys := strings.Join(keys, "")
	keyPair, err := provider.kvclient.Get(joinedKeys)
	if err != nil {
		log.Debugf("Cannot get key %s %s, setting default empty", joinedKeys, err)
		return []string{}
	} else if keyPair == nil {
		log.Debugf("Cannot get key %s, setting default %empty", joinedKeys)
		return []string{}
	}
	return strings.Split(string(keyPair.Value), ",")
}

func (provider *Kv) last(key string) string {
	splittedKey := strings.Split(key, "/")
	return splittedKey[len(splittedKey)-1]
}

func (provider *Kv) checkConstraints(keys ...string) bool {
	joinedKeys := strings.Join(keys, "")
	keyPair, err := provider.kvclient.Get(joinedKeys)

	value := ""
	if err == nil && keyPair != nil && keyPair.Value != nil {
		value = string(keyPair.Value)
	}

	constraintTags := strings.Split(value, ",")
	ok, failingConstraint := provider.MatchConstraints(constraintTags)
	if ok == false {
		if failingConstraint != nil {
			log.Debugf("Constraint %v not matching with following tags: %v", failingConstraint.String(), value)
		}
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
		}
		go p.start(powerCallback)
	}
}

func getResolvedAddresses(localHost, remoteHost *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", *remoteHost)
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			// fmt.Printf("%#v", string(buff[:n]))
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			fmt.Printf("newPacket : %v\n", newPacket)
			fmt.Printf("len(r) : %v\n", len(r))
			// if remainingBytes > 0 {
			// 	if remainingBytes <= n {
			// 		fmt.Println("1")
			// 		newPacket = true
			// 		msg = msg + string(r.next(remainingBytes))
			// 		fmt.Printf("2 Remaining bytes: %d\n", remainingBytes)
			// 		remainingBytes = n - remainingBytes
			// 		fmt.Printf("3 Remaining bytes: %d\n", remainingBytes)
			// 		// fmt.Println("msg: ", string(msg))
			// 	} else {
			// 		fmt.Println("2")
			// 		newPacket = false
			// 		msg = msg + string(r.next(remainingBytes))
			// 		remainingBytes = remainingBytes - n
			// 	}
			//
			// 	fmt.Printf("1 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
			// }
			fmt.Println("3")
		NewP:
			fmt.Println("4")
			if newPacket && len(r) > 4 && remainingBytes == 0 {
				fmt.Println("5")
				// remainingBytes = 0
				newPacket = false
				fmt.Printf("2 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
				var msg []byte
				t := r.byte()
				n = n - 1
				fmt.Println("t: ", string(t))
				switch t {
				case 'Q', 'B', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'P', 'p', 'S', 'X':
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					if remainingBytes < 4 {
						fmt.Errorf("ERROR: remainingBytes can't be less than 4 bytes if int32")
					} else {
						remainingBytes = remainingBytes - 4
						if remainingBytes > 0 {
							if remainingBytes <= n {
								newPacket = true
								msg = append(msg, r.next(remainingBytes)[:]...)
								remainingBytes = n - remainingBytes
								// fmt.Printf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
								// fmt.Printf("3 Remaining bytes: %d \tmsg: %v\n", remainingBytes, msg)
								fmt.Printf("3 Remaining bytes: %d \tmsg: ", remainingBytes)
								for _, v := range msg {
									fmt.Printf("'%v': '%s'  ", v, string(v))
								}
								fmt.Println(msg)
								fmt.Println(string(msg))
								goto NewP
							} else {
								newPacket = false
								msg = append(msg, r.next(remainingBytes)[:]...)
								remainingBytes = remainingBytes - n
								fmt.Printf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							}
						}
					}
					// case :
					// 	fmt.Println("TODO")
					// 	// c.rxReadyForQuery(r)
					// 	remainingBytes = r.int32()
					// 	remainingBytes = remainingBytes - 4
					// 	if remainingBytes > 0 {
					// 		if remainingBytes <= n {
					// 			newPacket = true
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = n - remainingBytes
					// 			fmt.Printf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
					// 			// fmt.Println(msg)
					// 			goto NewP
					// 		} else {
					// 			newPacket = false
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = remainingBytes - n
					// 			fmt.Printf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
					// 		}
					// 	}
					// case rowDescription:
					// case dataRow:
					// case bindComplete:
					// case commandComplete:
					// 	commandTag = CommandTag(r.readCString())
				default:
					fmt.Println("6")
					remainingBytes = 0
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// fmt.Printf("t: %#v\n", t)

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}
<commit_msg>Update<commit_after>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
		}
		go p.start(powerCallback)
	}
}

func getResolvedAddresses(localHost, remoteHost *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", *remoteHost)
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			// fmt.Printf("%#v", string(buff[:n]))
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			fmt.Printf("newPacket : %v\n", newPacket)
			fmt.Printf("len(r) : %v\n", len(r))
			// if remainingBytes > 0 {
			// 	if remainingBytes <= n {
			// 		fmt.Println("1")
			// 		newPacket = true
			// 		msg = msg + string(r.next(remainingBytes))
			// 		fmt.Printf("2 Remaining bytes: %d\n", remainingBytes)
			// 		remainingBytes = n - remainingBytes
			// 		fmt.Printf("3 Remaining bytes: %d\n", remainingBytes)
			// 		// fmt.Println("msg: ", string(msg))
			// 	} else {
			// 		fmt.Println("2")
			// 		newPacket = false
			// 		msg = msg + string(r.next(remainingBytes))
			// 		remainingBytes = remainingBytes - n
			// 	}
			//
			// 	fmt.Printf("1 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
			// }
			fmt.Println("3")
		NewP:
			fmt.Println("4")
			if newPacket && len(r) > 4 && remainingBytes == 0 {
				fmt.Println("5")
				// remainingBytes = 0
				newPacket = false
				fmt.Printf("2 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
				var msg []byte
				t := r.byte()
				n = n - 1
				fmt.Println("t: ", string(t))
				switch t {
				case 'Q', 'B', 'C', 'd', 'c', 'f', 'D', 'E', 'H', 'F', 'P', 'p', 'S', 'X':
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					if remainingBytes < 4 {
						fmt.Errorf("ERROR: remainingBytes can't be less than 4 bytes if int32")
					} else {
						remainingBytes = remainingBytes - 4
						if remainingBytes > 0 {
							if remainingBytes <= n {
								newPacket = true
								msg = append(msg, r.next(remainingBytes)[:]...)
								remainingBytes = n - remainingBytes
								// fmt.Printf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
								// fmt.Printf("3 Remaining bytes: %d \tmsg: %v\n", remainingBytes, msg)
								fmt.Printf("3 Remaining bytes: %d \tmsg: ", remainingBytes)
								for _, v := range msg {
									fmt.Printf("'%v': '%s'  ", v, string(v))
								}
								fmt.Println(msg)
								fmt.Println(string(msg))
								goto NewP
							} else {
								newPacket = false
								msg = append(msg, r.next(remainingBytes)[:]...)
								remainingBytes = remainingBytes - n
								fmt.Printf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
							}
						}
					}
					// case :
					// 	fmt.Println("TODO")
					// 	// c.rxReadyForQuery(r)
					// 	remainingBytes = r.int32()
					// 	remainingBytes = remainingBytes - 4
					// 	if remainingBytes > 0 {
					// 		if remainingBytes <= n {
					// 			newPacket = true
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = n - remainingBytes
					// 			fmt.Printf("3 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
					// 			// fmt.Println(msg)
					// 			goto NewP
					// 		} else {
					// 			newPacket = false
					// 			msg = msg + string(r.next(remainingBytes))
					// 			remainingBytes = remainingBytes - n
					// 			fmt.Printf("4 Remaining bytes: %d \tmsg: %s\n", remainingBytes, string(msg))
					// 		}
					// 	}
					// case rowDescription:
					// case dataRow:
					// case bindComplete:
					// case commandComplete:
					// 	commandTag = CommandTag(r.readCString())
				default:
					fmt.Println("6")
					remainingBytes = 0
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			} else {
				fmt.Println("7")
				remainingBytes = 0
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// fmt.Printf("t: %#v\n", t)

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
		}
		go p.start(powerCallback)
	}
}

func getResolvedAddresses(localHost, remoteHost *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", *remoteHost)
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			// fmt.Printf("%#v", string(buff[:n]))
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			if remainingBytes > 0 {
				if remainingBytes <= n {
					fmt.Println("1")
					newPacket = true
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = n - remainingBytes
					fmt.Println("msg: ", string(msg))
				} else {
					fmt.Println("2")
					newPacket = false
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = remainingBytes - n
				}
			}
			fmt.Println("3")
		NewP:
			fmt.Println("4")
			if newPacket && len(r) > 0 {
				fmt.Println("5")
				remainingBytes = 0
				newPacket = false
				msg = ""
				t := r.byte()
				n = n - 1
				fmt.Println("t: ", string(t))
				switch t {
				case query:
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					remainingBytes = remainingBytes - 4
					if remainingBytes > 0 {
						if remainingBytes <= n {
							newPacket = true
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = n - remainingBytes
							fmt.Println(msg)
							goto NewP
						} else {
							newPacket = false
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = remainingBytes - n
						}
						fmt.Printf("Remaining bytes: %d\n", remainingBytes)
					}
				// case rowDescription:
				// case dataRow:
				// case bindComplete:
				// case commandComplete:
				// 	commandTag = CommandTag(r.readCString())
				default:
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// fmt.Printf("t: %#v\n", t)

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}
<commit_msg>Update<commit_after>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
		}
		go p.start(powerCallback)
	}
}

func getResolvedAddresses(localHost, remoteHost *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", *remoteHost)
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			// fmt.Printf("%#v", string(buff[:n]))
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			if remainingBytes > 0 {
				if remainingBytes <= n {
					fmt.Println("1")
					newPacket = true
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = n - remainingBytes
					fmt.Println("msg: ", string(msg))
				} else {
					fmt.Println("2")
					newPacket = false
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = remainingBytes - n
				}
			}
			fmt.Println("3")
		NewP:
			fmt.Println("4")
			if newPacket && len(r) > 0 {
				fmt.Println("5")
				remainingBytes = 0
				newPacket = false
				fmt.Println("msg 0 : ", string(msg))
				msg = ""
				t := r.byte()
				n = n - 1
				fmt.Println("t: ", string(t))
				switch t {
				case query:
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					remainingBytes = remainingBytes - 4
					if remainingBytes > 0 {
						if remainingBytes <= n {
							newPacket = true
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = n - remainingBytes
							fmt.Println("msg: ", string(msg))
							// fmt.Println(msg)
							goto NewP
						} else {
							newPacket = false
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = remainingBytes - n
						}
						fmt.Printf("Remaining bytes: %d\n", remainingBytes)
					}
				// case rowDescription:
				// case dataRow:
				// case bindComplete:
				// case commandComplete:
				// 	commandTag = CommandTag(r.readCString())
				default:
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// fmt.Printf("t: %#v\n", t)

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}
<|endoftext|>