<commit_before>package scenarios

import (
	"crypto/rand"
	"os"
	"runtime"
	"sync"
	"testing"
	"time"

	"v2ray.com/core"
	"v2ray.com/core/app/log"
	"v2ray.com/core/app/proxyman"
	"v2ray.com/core/common/compare"
	clog "v2ray.com/core/common/log"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/serial"
	"v2ray.com/core/common/uuid"
	"v2ray.com/core/proxy/dokodemo"
	"v2ray.com/core/proxy/freedom"
	"v2ray.com/core/proxy/vmess"
	"v2ray.com/core/proxy/vmess/inbound"
	"v2ray.com/core/proxy/vmess/outbound"
	"v2ray.com/core/testing/servers/tcp"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/domainsocket"
	"v2ray.com/core/transport/internet/headers/http"
	"v2ray.com/core/transport/internet/quic"
	tcptransport "v2ray.com/core/transport/internet/tcp"
	. "v2ray.com/ext/assert"
)

func TestHttpConnectionHeader(t *testing.T) {
	assert := With(t)

	tcpServer := tcp.Server{
		MsgProcessor: xor,
	}
	dest, err := tcpServer.Start()
	assert(err, IsNil)
	defer tcpServer.Close()

	userID := protocol.NewID(uuid.New())
	serverPort := tcp.PickPort()
	serverConfig := &core.Config{
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(serverPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
					StreamSettings: &internet.StreamConfig{
						TransportSettings: []*internet.TransportConfig{
							{
								Protocol: internet.TransportProtocol_TCP,
								Settings: serial.ToTypedMessage(&tcptransport.Config{
									HeaderSettings: serial.ToTypedMessage(&http.Config{}),
								}),
							},
						},
					},
				}),
				ProxySettings: serial.ToTypedMessage(&inbound.Config{
					User: []*protocol.User{
						{
							Account: serial.ToTypedMessage(&vmess.Account{
								Id: userID.String(),
							}),
						},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
			},
		},
	}

	clientPort := tcp.PickPort()
	clientConfig := &core.Config{
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(clientPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
				}),
				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
					Address: net.NewIPOrDomain(dest.Address),
					Port:    uint32(dest.Port),
					NetworkList: &net.NetworkList{
						Network: []net.Network{net.Network_TCP},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&outbound.Config{
					Receiver: []*protocol.ServerEndpoint{
						{
							Address: net.NewIPOrDomain(net.LocalHostIP),
							Port:    uint32(serverPort),
							User: []*protocol.User{
								{
									Account: serial.ToTypedMessage(&vmess.Account{
										Id: userID.String(),
									}),
								},
							},
						},
					},
				}),
				SenderSettings: serial.ToTypedMessage(&proxyman.SenderConfig{
					StreamSettings: &internet.StreamConfig{
						TransportSettings: []*internet.TransportConfig{
							{
								Protocol: internet.TransportProtocol_TCP,
								Settings: serial.ToTypedMessage(&tcptransport.Config{
									HeaderSettings: serial.ToTypedMessage(&http.Config{}),
								}),
							},
						},
					},
				}),
			},
		},
	}

	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
	assert(err, IsNil)

	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   []byte{127, 0, 0, 1},
		Port: int(clientPort),
	})
	assert(err, IsNil)

	payload := "dokodemo request."
	nBytes, err := conn.Write([]byte(payload))
	assert(err, IsNil)
	assert(nBytes, Equals, len(payload))

	response := readFrom(conn, time.Second*2, len(payload))
	assert(response, Equals, xor([]byte(payload)))
	assert(conn.Close(), IsNil)

	CloseAllServers(servers)
}

func TestDomainSocket(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("Not supported on windows")
		return
	}
	assert := With(t)

	tcpServer := tcp.Server{
		MsgProcessor: xor,
	}
	dest, err := tcpServer.Start()
	assert(err, IsNil)
	defer tcpServer.Close()

	const dsPath = "/tmp/ds_scenario"
	os.Remove(dsPath)

	userID := protocol.NewID(uuid.New())
	serverPort := tcp.PickPort()
	serverConfig := &core.Config{
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(serverPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
					StreamSettings: &internet.StreamConfig{
						Protocol: internet.TransportProtocol_DomainSocket,
						TransportSettings: []*internet.TransportConfig{
							{
								Protocol: internet.TransportProtocol_DomainSocket,
								Settings: serial.ToTypedMessage(&domainsocket.Config{
									Path: dsPath,
								}),
							},
						},
					},
				}),
				ProxySettings: serial.ToTypedMessage(&inbound.Config{
					User: []*protocol.User{
						{
							Account: serial.ToTypedMessage(&vmess.Account{
								Id: userID.String(),
							}),
						},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
			},
		},
	}

	clientPort := tcp.PickPort()
	clientConfig := &core.Config{
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(clientPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
				}),
				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
					Address: net.NewIPOrDomain(dest.Address),
					Port:    uint32(dest.Port),
					NetworkList: &net.NetworkList{
						Network: []net.Network{net.Network_TCP},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&outbound.Config{
					Receiver: []*protocol.ServerEndpoint{
						{
							Address: net.NewIPOrDomain(net.LocalHostIP),
							Port:    uint32(serverPort),
							User: []*protocol.User{
								{
									Account: serial.ToTypedMessage(&vmess.Account{
										Id: userID.String(),
									}),
								},
							},
						},
					},
				}),
				SenderSettings: serial.ToTypedMessage(&proxyman.SenderConfig{
					StreamSettings: &internet.StreamConfig{
						Protocol: internet.TransportProtocol_DomainSocket,
						TransportSettings: []*internet.TransportConfig{
							{
								Protocol: internet.TransportProtocol_DomainSocket,
								Settings: serial.ToTypedMessage(&domainsocket.Config{
									Path: dsPath,
								}),
							},
						},
					},
				}),
			},
		},
	}

	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
	assert(err, IsNil)

	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   []byte{127, 0, 0, 1},
		Port: int(clientPort),
	})
	assert(err, IsNil)

	payload := "dokodemo request."
	nBytes, err := conn.Write([]byte(payload))
	assert(err, IsNil)
	assert(nBytes, Equals, len(payload))

	response := readFrom(conn, time.Second*2, len(payload))
	assert(response, Equals, xor([]byte(payload)))
	assert(conn.Close(), IsNil)

	CloseAllServers(servers)
}

func TestVMessQuic(t *testing.T) {
	assert := With(t)

	tcpServer := tcp.Server{
		MsgProcessor: xor,
	}
	dest, err := tcpServer.Start()
	assert(err, IsNil)
	defer tcpServer.Close()

	userID := protocol.NewID(uuid.New())
	serverPort := tcp.PickPort()
	serverConfig := &core.Config{
		App: []*serial.TypedMessage{
			serial.ToTypedMessage(&log.Config{
				ErrorLogLevel: clog.Severity_Debug,
				ErrorLogType:  log.LogType_Console,
			}),
		},
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(serverPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
					StreamSettings: &internet.StreamConfig{
						ProtocolName: "quic",
						TransportSettings: []*internet.TransportConfig{
							{
								ProtocolName: "quic",
								Settings:     serial.ToTypedMessage(&quic.Config{}),
							},
						},
					},
				}),
				ProxySettings: serial.ToTypedMessage(&inbound.Config{
					User: []*protocol.User{
						{
							Account: serial.ToTypedMessage(&vmess.Account{
								Id:      userID.String(),
								AlterId: 64,
							}),
						},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
			},
		},
	}

	clientPort := tcp.PickPort()
	clientConfig := &core.Config{
		App: []*serial.TypedMessage{
			serial.ToTypedMessage(&log.Config{
				ErrorLogLevel: clog.Severity_Debug,
				ErrorLogType:  log.LogType_Console,
			}),
		},
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(clientPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
				}),
				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
					Address: net.NewIPOrDomain(dest.Address),
					Port:    uint32(dest.Port),
					NetworkList: &net.NetworkList{
						Network: []net.Network{net.Network_TCP},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				SenderSettings: serial.ToTypedMessage(&proxyman.SenderConfig{
					StreamSettings: &internet.StreamConfig{
						ProtocolName: "quic",
						TransportSettings: []*internet.TransportConfig{
							{
								ProtocolName: "quic",
								Settings:     serial.ToTypedMessage(&quic.Config{}),
							},
						},
					},
				}),
				ProxySettings: serial.ToTypedMessage(&outbound.Config{
					Receiver: []*protocol.ServerEndpoint{
						{
							Address: net.NewIPOrDomain(net.LocalHostIP),
							Port:    uint32(serverPort),
							User: []*protocol.User{
								{
									Account: serial.ToTypedMessage(&vmess.Account{
										Id:      userID.String(),
										AlterId: 64,
										SecuritySettings: &protocol.SecurityConfig{
											Type: protocol.SecurityType_AES128_GCM,
										},
									}),
								},
							},
						},
					},
				}),
			},
		},
	}

	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
	if err != nil {
		t.Fatal("Failed to initialize all servers: ", err.Error())
	}
	defer CloseAllServers(servers)

	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
				IP:   []byte{127, 0, 0, 1},
				Port: int(clientPort),
			})
			assert(err, IsNil)
			defer conn.Close() // nolint: errcheck

			payload := make([]byte, 10240*1024)
			rand.Read(payload)

			nBytes, err := conn.Write([]byte(payload))
			assert(err, IsNil)
			assert(nBytes, Equals, len(payload))

			response := readFrom(conn, time.Second*40, 10240*1024)
			if err := compare.BytesEqualWithDetail(response, xor([]byte(payload))); err != nil {
				t.Error(err)
			}
		}()
	}
	wg.Wait()
}
<commit_msg>update settings<commit_after>package scenarios

import (
	"crypto/rand"
	"os"
	"runtime"
	"sync"
	"testing"
	"time"

	"v2ray.com/core/transport/internet/headers/wechat"

	"v2ray.com/core"
	"v2ray.com/core/app/log"
	"v2ray.com/core/app/proxyman"
	"v2ray.com/core/common/compare"
	clog "v2ray.com/core/common/log"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/serial"
	"v2ray.com/core/common/uuid"
	"v2ray.com/core/proxy/dokodemo"
	"v2ray.com/core/proxy/freedom"
	"v2ray.com/core/proxy/vmess"
	"v2ray.com/core/proxy/vmess/inbound"
	"v2ray.com/core/proxy/vmess/outbound"
	"v2ray.com/core/testing/servers/tcp"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/domainsocket"
	"v2ray.com/core/transport/internet/headers/http"
	"v2ray.com/core/transport/internet/quic"
	tcptransport "v2ray.com/core/transport/internet/tcp"
	. "v2ray.com/ext/assert"
)

func TestHttpConnectionHeader(t *testing.T) {
	assert := With(t)

	tcpServer := tcp.Server{
		MsgProcessor: xor,
	}
	dest, err := tcpServer.Start()
	assert(err, IsNil)
	defer tcpServer.Close()

	userID := protocol.NewID(uuid.New())
	serverPort := tcp.PickPort()
	serverConfig := &core.Config{
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(serverPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
					StreamSettings: &internet.StreamConfig{
						TransportSettings: []*internet.TransportConfig{
							{
								Protocol: internet.TransportProtocol_TCP,
								Settings: serial.ToTypedMessage(&tcptransport.Config{
									HeaderSettings: serial.ToTypedMessage(&http.Config{}),
								}),
							},
						},
					},
				}),
				ProxySettings: serial.ToTypedMessage(&inbound.Config{
					User: []*protocol.User{
						{
							Account: serial.ToTypedMessage(&vmess.Account{
								Id: userID.String(),
							}),
						},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
			},
		},
	}

	clientPort := tcp.PickPort()
	clientConfig := &core.Config{
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(clientPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
				}),
				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
					Address: net.NewIPOrDomain(dest.Address),
					Port:    uint32(dest.Port),
					NetworkList: &net.NetworkList{
						Network: []net.Network{net.Network_TCP},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&outbound.Config{
					Receiver: []*protocol.ServerEndpoint{
						{
							Address: net.NewIPOrDomain(net.LocalHostIP),
							Port:    uint32(serverPort),
							User: []*protocol.User{
								{
									Account: serial.ToTypedMessage(&vmess.Account{
										Id: userID.String(),
									}),
								},
							},
						},
					},
				}),
				SenderSettings: serial.ToTypedMessage(&proxyman.SenderConfig{
					StreamSettings: &internet.StreamConfig{
						TransportSettings: []*internet.TransportConfig{
							{
								Protocol: internet.TransportProtocol_TCP,
								Settings: serial.ToTypedMessage(&tcptransport.Config{
									HeaderSettings: serial.ToTypedMessage(&http.Config{}),
								}),
							},
						},
					},
				}),
			},
		},
	}

	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
	assert(err, IsNil)

	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   []byte{127, 0, 0, 1},
		Port: int(clientPort),
	})
	assert(err, IsNil)

	payload := "dokodemo request."
	nBytes, err := conn.Write([]byte(payload))
	assert(err, IsNil)
	assert(nBytes, Equals, len(payload))

	response := readFrom(conn, time.Second*2, len(payload))
	assert(response, Equals, xor([]byte(payload)))
	assert(conn.Close(), IsNil)

	CloseAllServers(servers)
}

func TestDomainSocket(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("Not supported on windows")
		return
	}
	assert := With(t)

	tcpServer := tcp.Server{
		MsgProcessor: xor,
	}
	dest, err := tcpServer.Start()
	assert(err, IsNil)
	defer tcpServer.Close()

	const dsPath = "/tmp/ds_scenario"
	os.Remove(dsPath)

	userID := protocol.NewID(uuid.New())
	serverPort := tcp.PickPort()
	serverConfig := &core.Config{
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(serverPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
					StreamSettings: &internet.StreamConfig{
						Protocol: internet.TransportProtocol_DomainSocket,
						TransportSettings: []*internet.TransportConfig{
							{
								Protocol: internet.TransportProtocol_DomainSocket,
								Settings: serial.ToTypedMessage(&domainsocket.Config{
									Path: dsPath,
								}),
							},
						},
					},
				}),
				ProxySettings: serial.ToTypedMessage(&inbound.Config{
					User: []*protocol.User{
						{
							Account: serial.ToTypedMessage(&vmess.Account{
								Id: userID.String(),
							}),
						},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
			},
		},
	}

	clientPort := tcp.PickPort()
	clientConfig := &core.Config{
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(clientPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
				}),
				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
					Address: net.NewIPOrDomain(dest.Address),
					Port:    uint32(dest.Port),
					NetworkList: &net.NetworkList{
						Network: []net.Network{net.Network_TCP},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&outbound.Config{
					Receiver: []*protocol.ServerEndpoint{
						{
							Address: net.NewIPOrDomain(net.LocalHostIP),
							Port:    uint32(serverPort),
							User: []*protocol.User{
								{
									Account: serial.ToTypedMessage(&vmess.Account{
										Id: userID.String(),
									}),
								},
							},
						},
					},
				}),
				SenderSettings: serial.ToTypedMessage(&proxyman.SenderConfig{
					StreamSettings: &internet.StreamConfig{
						Protocol: internet.TransportProtocol_DomainSocket,
						TransportSettings: []*internet.TransportConfig{
							{
								Protocol: internet.TransportProtocol_DomainSocket,
								Settings: serial.ToTypedMessage(&domainsocket.Config{
									Path: dsPath,
								}),
							},
						},
					},
				}),
			},
		},
	}

	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
	assert(err, IsNil)

	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   []byte{127, 0, 0, 1},
		Port: int(clientPort),
	})
	assert(err, IsNil)

	payload := "dokodemo request."
	nBytes, err := conn.Write([]byte(payload))
	assert(err, IsNil)
	assert(nBytes, Equals, len(payload))

	response := readFrom(conn, time.Second*2, len(payload))
	assert(response, Equals, xor([]byte(payload)))
	assert(conn.Close(), IsNil)

	CloseAllServers(servers)
}

func TestVMessQuic(t *testing.T) {
	assert := With(t)

	tcpServer := tcp.Server{
		MsgProcessor: xor,
	}
	dest, err := tcpServer.Start()
	assert(err, IsNil)
	defer tcpServer.Close()

	userID := protocol.NewID(uuid.New())
	serverPort := tcp.PickPort()
	serverConfig := &core.Config{
		App: []*serial.TypedMessage{
			serial.ToTypedMessage(&log.Config{
				ErrorLogLevel: clog.Severity_Debug,
				ErrorLogType:  log.LogType_Console,
			}),
		},
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(serverPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
					StreamSettings: &internet.StreamConfig{
						ProtocolName: "quic",
						TransportSettings: []*internet.TransportConfig{
							{
								ProtocolName: "quic",
								Settings: serial.ToTypedMessage(&quic.Config{
									Header: serial.ToTypedMessage(&wechat.VideoConfig{}),
								}),
							},
						},
					},
				}),
				ProxySettings: serial.ToTypedMessage(&inbound.Config{
					User: []*protocol.User{
						{
							Account: serial.ToTypedMessage(&vmess.Account{
								Id:      userID.String(),
								AlterId: 64,
							}),
						},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
			},
		},
	}

	clientPort := tcp.PickPort()
	clientConfig := &core.Config{
		App: []*serial.TypedMessage{
			serial.ToTypedMessage(&log.Config{
				ErrorLogLevel: clog.Severity_Debug,
				ErrorLogType:  log.LogType_Console,
			}),
		},
		Inbound: []*core.InboundHandlerConfig{
			{
				ReceiverSettings: serial.ToTypedMessage(&proxyman.ReceiverConfig{
					PortRange: net.SinglePortRange(clientPort),
					Listen:    net.NewIPOrDomain(net.LocalHostIP),
				}),
				ProxySettings: serial.ToTypedMessage(&dokodemo.Config{
					Address: net.NewIPOrDomain(dest.Address),
					Port:    uint32(dest.Port),
					NetworkList: &net.NetworkList{
						Network: []net.Network{net.Network_TCP},
					},
				}),
			},
		},
		Outbound: []*core.OutboundHandlerConfig{
			{
				SenderSettings: serial.ToTypedMessage(&proxyman.SenderConfig{
					StreamSettings: &internet.StreamConfig{
						ProtocolName: "quic",
						TransportSettings: []*internet.TransportConfig{
							{
								ProtocolName: "quic",
								Settings: serial.ToTypedMessage(&quic.Config{
									Header: serial.ToTypedMessage(&wechat.VideoConfig{}),
								}),
							},
						},
					},
				}),
				ProxySettings: serial.ToTypedMessage(&outbound.Config{
					Receiver: []*protocol.ServerEndpoint{
						{
							Address: net.NewIPOrDomain(net.LocalHostIP),
							Port:    uint32(serverPort),
							User: []*protocol.User{
								{
									Account: serial.ToTypedMessage(&vmess.Account{
										Id:      userID.String(),
										AlterId: 64,
										SecuritySettings: &protocol.SecurityConfig{
											Type: protocol.SecurityType_AES128_GCM,
										},
									}),
								},
							},
						},
					},
				}),
			},
		},
	}

	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
	if err != nil {
		t.Fatal("Failed to initialize all servers: ", err.Error())
	}
	defer CloseAllServers(servers)

	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
				IP:   []byte{127, 0, 0, 1},
				Port: int(clientPort),
			})
			assert(err, IsNil)
			defer conn.Close() // nolint: errcheck

			payload := make([]byte, 10240*1024)
			rand.Read(payload)

			nBytes, err := conn.Write([]byte(payload))
			assert(err, IsNil)
			assert(nBytes, Equals, len(payload))

			response := readFrom(conn, time.Second*40, 10240*1024)
			if err := compare.BytesEqualWithDetail(response, xor([]byte(payload))); err != nil {
				t.Error(err)
			}
		}()
	}
	wg.Wait()
}
<|endoftext|>


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

package time_test

import (
	"testing"
	. "time"
)

func TestTicker(t *testing.T) {
	const Count = 10
	Delta := 100 * Millisecond
	ticker := NewTicker(Delta)
	t0 := Now()
	for i := 0; i < Count; i++ {
		<-ticker.C
	}
	ticker.Stop()
	t1 := Now()
	dt := t1.Sub(t0)
	target := Delta * Count
	slop := target * 2 / 10
	if dt < target-slop || (!testing.Short() && dt > target+slop) {
		t.Fatalf("%d %s ticks took %s, expected [%s,%s]", Count, Delta, dt, target-slop, target+slop)
	}
	// Now test that the ticker stopped
	Sleep(2 * Delta)
	select {
	case <-ticker.C:
		t.Fatal("Ticker did not shut down")
	default:
		// ok
	}
}

// Issue 21874
func TestTickerStopWithDirectInitialization(t *testing.T) {
	c := make(chan Time)
	tk := &Ticker{C: c}
	tk.Stop()
}

// Test that a bug tearing down a ticker has been fixed. This routine should not deadlock.
func TestTeardown(t *testing.T) {
	Delta := 100 * Millisecond
	if testing.Short() {
		Delta = 20 * Millisecond
	}
	for i := 0; i < 3; i++ {
		ticker := NewTicker(Delta)
		<-ticker.C
		ticker.Stop()
	}
}

// Test the Tick convenience wrapper.
func TestTick(t *testing.T) {
	// Test that giving a negative duration returns nil.
	if got := Tick(-1); got != nil {
		t.Errorf("Tick(-1) = %v; want nil", got)
	}
}

// Test that NewTicker panics when given a duration less than zero.
func TestNewTickerLtZeroDuration(t *testing.T) {
	defer func() {
		if err := recover(); err == nil {
			t.Errorf("NewTicker(-1) should have panicked")
		}
	}()
	NewTicker(-1)
}

func BenchmarkTicker(b *testing.B) {
	benchmark(b, func(n int) {
		ticker := NewTicker(Nanosecond)
		for i := 0; i < n; i++ {
			<-ticker.C
		}
		ticker.Stop()
	})
}
<commit_msg>time: deflake TestTicker<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package time_test

import (
	"fmt"
	"testing"
	. "time"
)

func TestTicker(t *testing.T) {
	// We want to test that a ticker takes as much time as expected.
	// Since we don't want the test to run for too long, we don't
	// want to use lengthy times. This makes the test inherently flaky.
	// So only report an error if it fails five times in a row.

	const count = 10
	delta := 20 * Millisecond

	var errs []string
	logErrs := func() {
		for _, e := range errs {
			t.Log(e)
		}
	}

	for i := 0; i < 5; i++ {
		ticker := NewTicker(delta)
		t0 := Now()
		for i := 0; i < count; i++ {
			<-ticker.C
		}
		ticker.Stop()
		t1 := Now()
		dt := t1.Sub(t0)
		target := delta * count
		slop := target * 2 / 10
		if dt < target-slop || dt > target+slop {
			errs = append(errs, fmt.Sprintf("%d %s ticks took %s, expected [%s,%s]", count, delta, dt, target-slop, target+slop))
			continue
		}
		// Now test that the ticker stopped.
		Sleep(2 * delta)
		select {
		case <-ticker.C:
			errs = append(errs, "Ticker did not shut down")
			continue
		default:
			// ok
		}

		// Test passed, so all done.
		if len(errs) > 0 {
			t.Logf("saw %d errors, ignoring to avoid flakiness", len(errs))
			logErrs()
		}

		return
	}

	t.Errorf("saw %d errors", len(errs))
	logErrs()
}

// Issue 21874
func TestTickerStopWithDirectInitialization(t *testing.T) {
	c := make(chan Time)
	tk := &Ticker{C: c}
	tk.Stop()
}

// Test that a bug tearing down a ticker has been fixed. This routine should not deadlock.
func TestTeardown(t *testing.T) {
	Delta := 100 * Millisecond
	if testing.Short() {
		Delta = 20 * Millisecond
	}
	for i := 0; i < 3; i++ {
		ticker := NewTicker(Delta)
		<-ticker.C
		ticker.Stop()
	}
}

// Test the Tick convenience wrapper.
func TestTick(t *testing.T) {
	// Test that giving a negative duration returns nil.
	if got := Tick(-1); got != nil {
		t.Errorf("Tick(-1) = %v; want nil", got)
	}
}

// Test that NewTicker panics when given a duration less than zero.
func TestNewTickerLtZeroDuration(t *testing.T) {
	defer func() {
		if err := recover(); err == nil {
			t.Errorf("NewTicker(-1) should have panicked")
		}
	}()
	NewTicker(-1)
}

func BenchmarkTicker(b *testing.B) {
	benchmark(b, func(n int) {
		ticker := NewTicker(Nanosecond)
		for i := 0; i < n; i++ {
			<-ticker.C
		}
		ticker.Stop()
	})
}
<|endoftext|>


<commit_before>// package chunk encodes timeseries in chunks of data
// see devdocs/chunk-format.md for more information.
package chunk

import (
	"fmt"

	"github.com/grafana/metrictank/mdata/chunk/tsz"
	"github.com/grafana/metrictank/mdata/errors"
)

// Chunk is a chunk of data. not concurrency safe.
// last check that the methods are being called safely by Dieter on 20/11/2018
// checked: String, Push, Finish, Encode and properties Series, NumPoints, First
// for the most part, confirming serialized access is easy by tracking all callers/references.
// The main exception is the ChunkWriteRequest mechanism. any CWR created is processed
// asynchronously, and chunk properties will be read (the series.T0, chunk.Encode(), etc)
// But it can be proven that we only call NewChunkWriteRequest() on chunks that are no longer being modified.
type Chunk struct {
	Series    tsz.SeriesLong
	NumPoints uint32
	First     bool
}

func New(t0 uint32) *Chunk {
	return &Chunk{
		Series: *tsz.NewSeriesLong(t0),
	}
}

func NewFirst(t0 uint32) *Chunk {
	return &Chunk{
		Series: *tsz.NewSeriesLong(t0),
		First:  true,
	}
}

func (c *Chunk) String() string {
	return fmt.Sprintf("<chunk T0=%d, LastTs=%d, NumPoints=%d, First=%t, Closed=%t>", c.Series.T0, c.Series.T, c.NumPoints, c.First, c.Series.Finished)
}

func (c *Chunk) Push(t uint32, v float64) error {
	if t == c.Series.T {
		return errors.ErrMetricNewValueForTimestamp
	} else if t < c.Series.T {
		return errors.ErrMetricTooOld
	}
	c.Series.Push(t, v)
	c.NumPoints++
	return nil
}

func (c *Chunk) Finish() {
	c.Series.Finish()
}

// Encode encodes the chunk
// note: chunks don't know their own span, the caller/owner manages that,
// so for formats that encode it, it needs to be passed in.
// the returned value contains no references to the chunk. data is copied.
func (c *Chunk) Encode(span uint32) []byte {
	return encode(span, FormatGoTszLongWithSpan, c.Series.Bytes())
}
<commit_msg>update Chunk safe access checks<commit_after>// package chunk encodes timeseries in chunks of data
// see devdocs/chunk-format.md for more information.
package chunk

import (
	"fmt"

	"github.com/grafana/metrictank/mdata/chunk/tsz"
	"github.com/grafana/metrictank/mdata/errors"
)

// Chunk is a chunk of data. not concurrency safe.
// last check that all properties and methods are being accessed safely by Dieter on 2021-10-06
// for the most part, confirming serialized access is easy by tracking all callers/references.
type Chunk struct {
	Series    tsz.SeriesLong
	NumPoints uint32
	First     bool
}

func New(t0 uint32) *Chunk {
	return &Chunk{
		Series: *tsz.NewSeriesLong(t0),
	}
}

func NewFirst(t0 uint32) *Chunk {
	return &Chunk{
		Series: *tsz.NewSeriesLong(t0),
		First:  true,
	}
}

func (c *Chunk) String() string {
	return fmt.Sprintf("<chunk T0=%d, LastTs=%d, NumPoints=%d, First=%t, Closed=%t>", c.Series.T0, c.Series.T, c.NumPoints, c.First, c.Series.Finished)
}

func (c *Chunk) Push(t uint32, v float64) error {
	if t == c.Series.T {
		return errors.ErrMetricNewValueForTimestamp
	} else if t < c.Series.T {
		return errors.ErrMetricTooOld
	}
	c.Series.Push(t, v)
	c.NumPoints++
	return nil
}

func (c *Chunk) Finish() {
	c.Series.Finish()
}

// Encode encodes the chunk
// note: chunks don't know their own span, the caller/owner manages that,
// so for formats that encode it, it needs to be passed in.
// the returned value contains no references to the chunk. data is copied.
func (c *Chunk) Encode(span uint32) []byte {
	return encode(span, FormatGoTszLongWithSpan, c.Series.Bytes())
}
<|endoftext|>


<commit_before>/**
 * MediaType
 *
 * Copyright © 2014 Trevor N. Suarez (Rican7)
 */

// This file provides provides a mutable implementation of the MediaType interface

package mediatype

import (
	"mime"
	"strings"
)

/**
 * Types
 */

// A mutable struct defining the components of a Media Type
type MediaTypeMutable struct {
	Main   string
	Tree   []string
	Sub    string
	Suf    string
	Params map[string]string
}

// Return a New instance of a MediaType struct
func NewMutable() MediaType {
	return &MediaTypeMutable{}
}

// Get the "main" (top-level) type as a string
func (m *MediaTypeMutable) MainType() string {
	return m.Main
}

// Get the "sub" type as a string
func (m *MediaTypeMutable) SubType() string {
	return m.Sub
}

// Get the split "sub" type as an array of strings split by the namespace separator
func (m *MediaTypeMutable) Trees() []string {
	return m.Tree
}

// Get the prefix of the type's trees
func (m *MediaTypeMutable) Prefix() string {
	if 0 < len(m.Tree) {
		return m.Tree[0]
	}

	return ""
}

// Get the "suffix" of the type as a string
func (m *MediaTypeMutable) Suffix() string {
	return m.Suf
}

// Get the defined parameters of the media type
func (m *MediaTypeMutable) Parameters() map[string]string {
	return m.Params
}

// Get the normalized type and sub-type as a string
func (m *MediaTypeMutable) FullType() string {
	var fullType string

	fullType += m.Main + MainSubSplitCharacter

	if 0 < len(m.Tree) {
		fullType += strings.Join(m.Tree, TreeSeparatorCharacter)
		fullType += TreeSeparatorCharacter
	}

	fullType += m.Sub

	if "" != m.Suf {
		fullType += SuffixCharacter + m.Suf
	}

	return fullType
}

// Get a string representation conforming to RFC 2045 and RFC 2616
func (m *MediaTypeMutable) String() string {
	return mime.FormatMediaType(m.FullType(), m.Params)
}

// Split the full type string into parts and assign those values to our struct
func splitTypes(fullType string) *MediaTypeMutable {
	var mt MediaTypeMutable

	// Split the main/sub types
	mainSubSplit := strings.Split(fullType, MainSubSplitCharacter)

	mt.Main = mainSubSplit[0]

	// If we got more than one part, we must have a sub-type
	if 1 < len(mainSubSplit) {
		// Split the remaining main/sub split from a possible suffix
		subSuffixSplit := strings.Split(mainSubSplit[1], SuffixCharacter)

		// If we got more than one part, we must have a suffix
		if 1 < len(subSuffixSplit) {
			mt.Suf = subSuffixSplit[1]
		}

		// Split the sub-type split into the possibly different trees
		treeSubSplit := strings.Split(subSuffixSplit[0], TreeSeparatorCharacter)
		treeSubSplitLength := len(treeSubSplit)

		mt.Sub = treeSubSplit[treeSubSplitLength-1]

		// If we got more than one part, we must have tree definitions
		if 1 < treeSubSplitLength {
			mt.Tree = treeSubSplit[0 : treeSubSplitLength-1]
		}
	}

	// Build from the raw
	return &mt
}
<commit_msg>Fixing the `FullType()` method based on the tests<commit_after>/**
 * MediaType
 *
 * Copyright © 2014 Trevor N. Suarez (Rican7)
 */

// This file provides provides a mutable implementation of the MediaType interface

package mediatype

import (
	"mime"
	"strings"
)

/**
 * Types
 */

// A mutable struct defining the components of a Media Type
type MediaTypeMutable struct {
	Main   string
	Tree   []string
	Sub    string
	Suf    string
	Params map[string]string
}

// Return a New instance of a MediaType struct
func NewMutable() MediaType {
	return &MediaTypeMutable{}
}

// Get the "main" (top-level) type as a string
func (m *MediaTypeMutable) MainType() string {
	return m.Main
}

// Get the "sub" type as a string
func (m *MediaTypeMutable) SubType() string {
	return m.Sub
}

// Get the split "sub" type as an array of strings split by the namespace separator
func (m *MediaTypeMutable) Trees() []string {
	return m.Tree
}

// Get the prefix of the type's trees
func (m *MediaTypeMutable) Prefix() string {
	if 0 < len(m.Tree) {
		return m.Tree[0]
	}

	return ""
}

// Get the "suffix" of the type as a string
func (m *MediaTypeMutable) Suffix() string {
	return m.Suf
}

// Get the defined parameters of the media type
func (m *MediaTypeMutable) Parameters() map[string]string {
	return m.Params
}

// Get the normalized full type as a string
func (m *MediaTypeMutable) FullType() string {
	var fullType string

	fullType += m.Main

	// If there are other pieces
	if 0 < len(m.Tree) || "" != m.Sub || "" != m.Suf {
		fullType += MainSubSplitCharacter
	}

	if 0 < len(m.Tree) {
		fullType += strings.Join(m.Tree, TreeSeparatorCharacter)
		fullType += TreeSeparatorCharacter
	}

	fullType += m.Sub

	if "" != m.Suf {
		fullType += SuffixCharacter + m.Suf
	}

	return fullType
}

// Get a string representation conforming to RFC 2045 and RFC 2616
func (m *MediaTypeMutable) String() string {
	return mime.FormatMediaType(m.FullType(), m.Params)
}

// Split the full type string into parts and assign those values to our struct
func splitTypes(fullType string) *MediaTypeMutable {
	var mt MediaTypeMutable

	// Split the main/sub types
	mainSubSplit := strings.Split(fullType, MainSubSplitCharacter)

	mt.Main = mainSubSplit[0]

	// If we got more than one part, we must have a sub-type
	if 1 < len(mainSubSplit) {
		// Split the remaining main/sub split from a possible suffix
		subSuffixSplit := strings.Split(mainSubSplit[1], SuffixCharacter)

		// If we got more than one part, we must have a suffix
		if 1 < len(subSuffixSplit) {
			mt.Suf = subSuffixSplit[1]
		}

		// Split the sub-type split into the possibly different trees
		treeSubSplit := strings.Split(subSuffixSplit[0], TreeSeparatorCharacter)
		treeSubSplitLength := len(treeSubSplit)

		mt.Sub = treeSubSplit[treeSubSplitLength-1]

		// If we got more than one part, we must have tree definitions
		if 1 < treeSubSplitLength {
			mt.Tree = treeSubSplit[0 : treeSubSplitLength-1]
		}
	}

	// Build from the raw
	return &mt
}
<|endoftext|>


<commit_before>package mains

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/zetamatta/nyagos/dos"
	. "github.com/zetamatta/nyagos/ifdbg"
	"github.com/zetamatta/nyagos/lua"
	"github.com/zetamatta/nyagos/shell"
)

func versionOrStamp() string {
	if Version != "" {
		return Version
	} else {
		return "v" + Stamp
	}
}

type InterpreterT interface {
	Interpret(string) (int, error)
}

func loadScripts(it *shell.Cmd, L lua.Lua) error {
	exeName, exeNameErr := os.Executable()
	if exeNameErr != nil {
		fmt.Fprintln(os.Stderr, exeNameErr)
	}
	exeFolder := filepath.Dir(exeName)

	if !silentmode {
		fmt.Printf("Nihongo Yet Another GOing Shell %s-%s by %s & %s\n",
			versionOrStamp(),
			runtime.GOARCH,
			runtime.Version(),
			"Lua 5.3")
		fmt.Println("(c) 2014-2017 NYAOS.ORG <http://www.nyaos.org>")
	}

	nyagos_d := filepath.Join(exeFolder, "nyagos.d")
	nyagos_d_fd, nyagos_d_err := os.Open(nyagos_d)
	if nyagos_d_err == nil {
		defer nyagos_d_fd.Close()
		finfos, err := nyagos_d_fd.Readdir(-1)
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
		} else {
			for _, finfo1 := range finfos {
				name1 := finfo1.Name()
				if !strings.HasSuffix(strings.ToLower(name1), ".lua") {
					continue
				}
				path1 := filepath.Join(nyagos_d, name1)
				if DBG {
					println("load real ", path1)
				}
				if err := L.Source(path1); err != nil {
					fmt.Fprintf(os.Stderr, "%s: %s\n", name1, err.Error())
				}
			}
		}
	}
	if _, err := runLua(it, L, filepath.Join(exeFolder, ".nyagos")); err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
	}
	barNyagos(it, exeFolder, L)
	if err := dotNyagos(it, L); err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
	}
	barNyagos(it, dos.GetHome(), L)
	return nil
}

func runLua(it *shell.Cmd, L lua.Lua, fname string) ([]byte, error) {
	_, err := os.Stat(fname)
	if err != nil {
		if os.IsNotExist(err) {
			// println("pass " + fname + " (not exists)")
			return []byte{}, nil
		} else {
			return nil, err
		}
	}
	if _, err := L.LoadFile(fname, "bt"); err != nil {
		return nil, err
	}
	chank := L.Dump()
	if err := NyagosCallLua(L, it, 0, 0); err != nil {
		return nil, err
	}
	// println("Run: " + fname)
	return chank, nil
}

func dotNyagos(it *shell.Cmd, L lua.Lua) error {
	dot_nyagos := filepath.Join(dos.GetHome(), ".nyagos")
	dotStat, err := os.Stat(dot_nyagos)
	if err != nil {
		return nil
	}
	cachePath := filepath.Join(AppDataDir(), runtime.GOARCH+".nyagos.luac")
	cacheStat, err := os.Stat(cachePath)
	if err == nil && !dotStat.ModTime().After(cacheStat.ModTime()) {
		_, err = runLua(it, L, cachePath)
		return err
	}
	chank, err := runLua(it, L, dot_nyagos)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(cachePath, chank, os.FileMode(0644))
}

func barNyagos(it InterpreterT, folder string, L lua.Lua) {
	bar_nyagos := filepath.Join(folder, "_nyagos")
	fd, fd_err := os.Open(bar_nyagos)
	if fd_err != nil {
		return
	}
	defer fd.Close()
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		text := scanner.Text()
		text = doLuaFilter(L, text)
		_, err := it.Interpret(text)
		if err != nil {
			fmt.Fprint(os.Stderr, err.Error())
		}
	}
}
<commit_msg>(internal) make `barNyagos()` independ on Lua<commit_after>package mains

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/zetamatta/nyagos/dos"
	. "github.com/zetamatta/nyagos/ifdbg"
	"github.com/zetamatta/nyagos/lua"
	"github.com/zetamatta/nyagos/shell"
)

func versionOrStamp() string {
	if Version != "" {
		return Version
	} else {
		return "v" + Stamp
	}
}

type InterpreterT interface {
	Interpret(string) (int, error)
}

func loadScripts(it *shell.Cmd, L lua.Lua) error {
	exeName, exeNameErr := os.Executable()
	if exeNameErr != nil {
		fmt.Fprintln(os.Stderr, exeNameErr)
	}
	exeFolder := filepath.Dir(exeName)

	if !silentmode {
		fmt.Printf("Nihongo Yet Another GOing Shell %s-%s by %s & %s\n",
			versionOrStamp(),
			runtime.GOARCH,
			runtime.Version(),
			"Lua 5.3")
		fmt.Println("(c) 2014-2017 NYAOS.ORG <http://www.nyaos.org>")
	}

	nyagos_d := filepath.Join(exeFolder, "nyagos.d")
	nyagos_d_fd, nyagos_d_err := os.Open(nyagos_d)
	if nyagos_d_err == nil {
		defer nyagos_d_fd.Close()
		finfos, err := nyagos_d_fd.Readdir(-1)
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
		} else {
			for _, finfo1 := range finfos {
				name1 := finfo1.Name()
				if !strings.HasSuffix(strings.ToLower(name1), ".lua") {
					continue
				}
				path1 := filepath.Join(nyagos_d, name1)
				if DBG {
					println("load real ", path1)
				}
				if err := L.Source(path1); err != nil {
					fmt.Fprintf(os.Stderr, "%s: %s\n", name1, err.Error())
				}
			}
		}
	}
	if _, err := runLua(it, L, filepath.Join(exeFolder, ".nyagos")); err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
	}
	filter := func(text string) string {
		return doLuaFilter(L, text)
	}
	barNyagos(it, exeFolder, filter)
	if err := dotNyagos(it, L); err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
	}
	barNyagos(it, dos.GetHome(), filter)
	return nil
}

func runLua(it *shell.Cmd, L lua.Lua, fname string) ([]byte, error) {
	_, err := os.Stat(fname)
	if err != nil {
		if os.IsNotExist(err) {
			// println("pass " + fname + " (not exists)")
			return []byte{}, nil
		} else {
			return nil, err
		}
	}
	if _, err := L.LoadFile(fname, "bt"); err != nil {
		return nil, err
	}
	chank := L.Dump()
	if err := NyagosCallLua(L, it, 0, 0); err != nil {
		return nil, err
	}
	// println("Run: " + fname)
	return chank, nil
}

func dotNyagos(it *shell.Cmd, L lua.Lua) error {
	dot_nyagos := filepath.Join(dos.GetHome(), ".nyagos")
	dotStat, err := os.Stat(dot_nyagos)
	if err != nil {
		return nil
	}
	cachePath := filepath.Join(AppDataDir(), runtime.GOARCH+".nyagos.luac")
	cacheStat, err := os.Stat(cachePath)
	if err == nil && !dotStat.ModTime().After(cacheStat.ModTime()) {
		_, err = runLua(it, L, cachePath)
		return err
	}
	chank, err := runLua(it, L, dot_nyagos)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(cachePath, chank, os.FileMode(0644))
}

func barNyagos(it InterpreterT, folder string, filter func(string) string) {
	bar_nyagos := filepath.Join(folder, "_nyagos")
	fd, fd_err := os.Open(bar_nyagos)
	if fd_err != nil {
		return
	}
	defer fd.Close()
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		_, err := it.Interpret(filter(scanner.Text()))
		if err != nil {
			fmt.Fprint(os.Stderr, err.Error())
		}
	}
}
<|endoftext|>


<commit_before>package markov

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/jmhodges/levigo"
	"log"
	"math/rand"
	"os"
	"strings"
)

var (
	MaxPhraseLength = 50
)

func tokenize(order int, sentence string) ([][]string, error) {
	var result [][]string

	words := strings.Fields(sentence)
	if len(words) < order {
		return result, fmt.Errorf("Sentence too short for order %d\n", order)
	}
	words = append(words, "\n")

	for i := 0; i < len(words)-order; i++ {
		ngram := words[i : i+order+1]
		result = append(result, ngram)
	}
	return result, nil
}

func MakeKey(ngram []string) ([]byte, error) {
	return json.Marshal(ngram)
}

type MarkovDB struct {
	Order int
	Db    *levigo.DB
}

func NewMarkovDB(order int, dbfile string) (*MarkovDB, error) {
	opts := levigo.NewOptions()
	opts.SetCache(levigo.NewLRUCache(3 << 29))
	opts.SetCreateIfMissing(true)
	db, err := levigo.Open(dbfile, opts)
	if err != nil {
		return nil, err
	}

	mdb := &MarkovDB{
		Order: order,
		Db:    db,
	}

	return mdb, nil
}

func (mdb *MarkovDB) ReadSentence(sentence string) {
	tokens, err := tokenize(mdb.Order, sentence)
	if err != nil {
		log.Println("ReadSentence error:", err)
		return
	}
	for _, token := range tokens {
		ngram := token[0 : len(token)-1]
		follow := token[len(token)-1]
		key, err := MakeKey(ngram)
		if err != nil {
			log.Println("ReadSentence error:", err)
			return
		}
		mdb.Put(key, follow)
	}
}

func (mdb *MarkovDB) Put(key []byte, value string) error {
	ro := levigo.NewReadOptions()
	wo := levigo.NewWriteOptions()
	defer ro.Close()
	defer wo.Close()

	// get existing value for key
	jsonwords, err := mdb.Db.Get(ro, key)
	if err != nil {
		return err
	}

	// deserialize existing value, if any
	var words []string
	if jsonwords != nil && string(jsonwords) != "" {
		err = json.Unmarshal(jsonwords, &words)
		if err != nil {
			return err
		}
	}

	// append the new word to the value and serialize the result
	for _, word := range words {
		if word == value {
			return nil
		}
	}
	words = append(words, value)
	newwords, err := json.Marshal(words)
	if err != nil {
		return err
	}

	// write the new value to the db
	err = mdb.Db.Put(wo, key, newwords)
	return err
}

func (mdb *MarkovDB) Generate(seed string) string {
	ro := levigo.NewReadOptions()
	defer ro.Close()

	var phrases []string

	tokens, err := tokenize(mdb.Order, seed)
	if err != nil {
		log.Println("Generate() error inside tokenize:", err)
		return ""
	}
	for _, token := range tokens {
		ngram := token[0 : len(token)-1]
		// fmt.Printf("ngram = %q\n", ngram)

		phrase := mdb.Goo(ngram)
		if phrase == seed {
			continue
		}
		phrases = append(phrases, phrase)
		// fmt.Printf("* phrase %d: %s\n", i, phrase)
	}

	var result string
	for _, phrase := range phrases {
		if len(phrase) > len(result) {
			result = phrase
		}
	}

	// fmt.Printf("%s\n", result)

	return result
}

func (mdb *MarkovDB) Goo(ngramKey []string) string {
	key, err := MakeKey(ngramKey)
	if err != nil {
		log.Println("Goo error (MakeKey):", err)
		return ""
	}
	var result []string = make([]string, len(ngramKey))
	copy(result, ngramKey)

	// fmt.Printf("result (1): %q (%q)\n", result, ngramKey)

	for i := 0; i < MaxPhraseLength; i++ {
		followWord, err := mdb.GetRandom(key)
		if err != nil || followWord == "\n" {
			// fmt.Printf("no follow for %s\n", key)
			break
		}
		if followWord == "\n" {
			break
		}
		result = append(result, followWord)
		ngramKey = append(ngramKey[1:], followWord)
		key, err = MakeKey(ngramKey)
		if err != nil {
			break
		}
	}
	return strings.Join(result, " ")
}

func (mdb *MarkovDB) GetRandom(key []byte) (string, error) {
	ro := levigo.NewReadOptions()
	defer ro.Close()

	jsondata, err := mdb.Db.Get(ro, key)
	if err != nil {
		return "", nil
	}

	var follows []string
	if jsondata == nil || string(jsondata) == "" {
		return "", fmt.Errorf("no value for this key: %s\n", key)
	}
	err = json.Unmarshal(jsondata, &follows)
	if err != nil {
		return "", err
	}

	word := follows[rand.Intn(len(follows))]

	// fmt.Printf("random for %q: %q\n", key, word)

	return word, nil
}

func (mdb *MarkovDB) Close() {
	mdb.Db.Close()
}

func TestMarkov() {
	mdb, err := NewMarkovDB(2, "petodb")
	if err != nil {
		log.Fatal(err)
	}
	defer mdb.Close()

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("> ")
		text, _ := reader.ReadString('\n')
		if strings.HasPrefix(text, "quit") {
			break
		}
		// mdb.ReadSentence(text)
		mdb.Generate(text)
	}
}

func ReadStdin(dbpath string, order int) {
	mdb, err := NewMarkovDB(order, dbpath)
	if err != nil {
		log.Fatal(err)
	}
	defer mdb.Close()

	i := 1
	var buf string

	reader := bufio.NewReader(os.Stdin)
	for {
		text, err := reader.ReadString('\n')
		if err != nil {
			log.Print(err)
			break
		}

		text = strings.TrimRight(text, "\n")
		if text == "" {
			text = "\n"
		}

		if i%2 == 0 {
			fmt.Printf("Put %s -> %q\n", buf, text)
			mdb.Put([]byte(buf), text)
		} else {
			buf = text
		}

		i++
	}
}

// To read from a IRC log file with a format like:
//
//    Dec 19 15:24:41 <user>    hello world!
//
// You can cleanup the log running:
//
// awk '$4 ~ /^</ { print }' < irc.log | cut -d '>' -f 2- | sed 's/^_TAB_//'
//
// Where _TAB_ is 'C-v TAB' from a bash shell.

func ReadFile(dbpath, filename string, order int) error {
	mdb, err := NewMarkovDB(order, dbpath)
	if err != nil {
		log.Fatal(err)
	}
	defer mdb.Close()

	var reader *bufio.Reader

	if filename == "-" {
		reader = bufio.NewReader(os.Stdin)
	} else {
		file, err := os.Open(filename)
		if err != nil {
			return err
		}
		reader = bufio.NewReader(file)
	}

	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			return err
		}
		line = strings.TrimRight(line, "\r\n")
		mdb.ReadSentence(line)
	}
}
<commit_msg>do not print errors for short markov sentences<commit_after>package markov

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/jmhodges/levigo"
	"log"
	"math/rand"
	"os"
	"strings"
)

var (
	MaxPhraseLength = 50
)

func tokenize(order int, sentence string) ([][]string, error) {
	var result [][]string

	words := strings.Fields(sentence)
	if len(words) < order {
		return result, fmt.Errorf("Sentence too short for order %d\n", order)
	}
	words = append(words, "\n")

	for i := 0; i < len(words)-order; i++ {
		ngram := words[i : i+order+1]
		result = append(result, ngram)
	}
	return result, nil
}

func MakeKey(ngram []string) ([]byte, error) {
	return json.Marshal(ngram)
}

type MarkovDB struct {
	Order int
	Db    *levigo.DB
}

func NewMarkovDB(order int, dbfile string) (*MarkovDB, error) {
	opts := levigo.NewOptions()
	opts.SetCache(levigo.NewLRUCache(3 << 29))
	opts.SetCreateIfMissing(true)
	db, err := levigo.Open(dbfile, opts)
	if err != nil {
		return nil, err
	}

	mdb := &MarkovDB{
		Order: order,
		Db:    db,
	}

	return mdb, nil
}

func (mdb *MarkovDB) ReadSentence(sentence string) {
	tokens, err := tokenize(mdb.Order, sentence)
	if err != nil {
		log.Println("ReadSentence error:", err)
		return
	}
	for _, token := range tokens {
		ngram := token[0 : len(token)-1]
		follow := token[len(token)-1]
		key, err := MakeKey(ngram)
		if err != nil {
			log.Println("ReadSentence error:", err)
			return
		}
		mdb.Put(key, follow)
	}
}

func (mdb *MarkovDB) Put(key []byte, value string) error {
	ro := levigo.NewReadOptions()
	wo := levigo.NewWriteOptions()
	defer ro.Close()
	defer wo.Close()

	// get existing value for key
	jsonwords, err := mdb.Db.Get(ro, key)
	if err != nil {
		return err
	}

	// deserialize existing value, if any
	var words []string
	if jsonwords != nil && string(jsonwords) != "" {
		err = json.Unmarshal(jsonwords, &words)
		if err != nil {
			return err
		}
	}

	// append the new word to the value and serialize the result
	for _, word := range words {
		if word == value {
			return nil
		}
	}
	words = append(words, value)
	newwords, err := json.Marshal(words)
	if err != nil {
		return err
	}

	// write the new value to the db
	err = mdb.Db.Put(wo, key, newwords)
	return err
}

func (mdb *MarkovDB) Generate(seed string) string {
	ro := levigo.NewReadOptions()
	defer ro.Close()

	var phrases []string

	tokens, err := tokenize(mdb.Order, seed)
	if err != nil {
		return ""
	}
	for _, token := range tokens {
		ngram := token[0 : len(token)-1]
		// fmt.Printf("ngram = %q\n", ngram)

		phrase := mdb.Goo(ngram)
		if phrase == seed {
			continue
		}
		phrases = append(phrases, phrase)
		// fmt.Printf("* phrase %d: %s\n", i, phrase)
	}

	var result string
	for _, phrase := range phrases {
		if len(phrase) > len(result) {
			result = phrase
		}
	}

	// fmt.Printf("%s\n", result)

	return result
}

func (mdb *MarkovDB) Goo(ngramKey []string) string {
	key, err := MakeKey(ngramKey)
	if err != nil {
		log.Println("Goo error (MakeKey):", err)
		return ""
	}
	var result []string = make([]string, len(ngramKey))
	copy(result, ngramKey)

	// fmt.Printf("result (1): %q (%q)\n", result, ngramKey)

	for i := 0; i < MaxPhraseLength; i++ {
		followWord, err := mdb.GetRandom(key)
		if err != nil || followWord == "\n" {
			// fmt.Printf("no follow for %s\n", key)
			break
		}
		if followWord == "\n" {
			break
		}
		result = append(result, followWord)
		ngramKey = append(ngramKey[1:], followWord)
		key, err = MakeKey(ngramKey)
		if err != nil {
			break
		}
	}
	return strings.Join(result, " ")
}

func (mdb *MarkovDB) GetRandom(key []byte) (string, error) {
	ro := levigo.NewReadOptions()
	defer ro.Close()

	jsondata, err := mdb.Db.Get(ro, key)
	if err != nil {
		return "", nil
	}

	var follows []string
	if jsondata == nil || string(jsondata) == "" {
		return "", fmt.Errorf("no value for this key: %s\n", key)
	}
	err = json.Unmarshal(jsondata, &follows)
	if err != nil {
		return "", err
	}

	word := follows[rand.Intn(len(follows))]

	// fmt.Printf("random for %q: %q\n", key, word)

	return word, nil
}

func (mdb *MarkovDB) Close() {
	mdb.Db.Close()
}

func TestMarkov() {
	mdb, err := NewMarkovDB(2, "petodb")
	if err != nil {
		log.Fatal(err)
	}
	defer mdb.Close()

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("> ")
		text, _ := reader.ReadString('\n')
		if strings.HasPrefix(text, "quit") {
			break
		}
		// mdb.ReadSentence(text)
		mdb.Generate(text)
	}
}

func ReadStdin(dbpath string, order int) {
	mdb, err := NewMarkovDB(order, dbpath)
	if err != nil {
		log.Fatal(err)
	}
	defer mdb.Close()

	i := 1
	var buf string

	reader := bufio.NewReader(os.Stdin)
	for {
		text, err := reader.ReadString('\n')
		if err != nil {
			log.Print(err)
			break
		}

		text = strings.TrimRight(text, "\n")
		if text == "" {
			text = "\n"
		}

		if i%2 == 0 {
			fmt.Printf("Put %s -> %q\n", buf, text)
			mdb.Put([]byte(buf), text)
		} else {
			buf = text
		}

		i++
	}
}

// To read from a IRC log file with a format like:
//
//    Dec 19 15:24:41 <user>    hello world!
//
// You can cleanup the log running:
//
// awk '$4 ~ /^</ { print }' < irc.log | cut -d '>' -f 2- | sed 's/^_TAB_//'
//
// Where _TAB_ is 'C-v TAB' from a bash shell.

func ReadFile(dbpath, filename string, order int) error {
	mdb, err := NewMarkovDB(order, dbpath)
	if err != nil {
		log.Fatal(err)
	}
	defer mdb.Close()

	var reader *bufio.Reader

	if filename == "-" {
		reader = bufio.NewReader(os.Stdin)
	} else {
		file, err := os.Open(filename)
		if err != nil {
			return err
		}
		reader = bufio.NewReader(file)
	}

	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			return err
		}
		line = strings.TrimRight(line, "\r\n")
		mdb.ReadSentence(line)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"syscall"

	"github.com/dstotijn/go-notion"
	"github.com/pelletier/go-toml"
)

func main() {
	if err := realmain(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

type Config struct {
	Currencies map[string]string
	Funds      map[string]string
}

func realmain() error {
	var (
		flagConfig = flag.String("c", "config.toml", "TOML configuration path")
	)
	flag.Parse()

	bytes, err := ioutil.ReadFile(*flagConfig)
	if err != nil {
		return err
	}

	var cfg Config
	if err := toml.Unmarshal(bytes, &cfg); err != nil {
		return err
	}

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	secret := os.Getenv("NOTION_YT_SECRET")
	if secret == "" {
		return fmt.Errorf("NOTION_YET_SECRET environment variable is not set")
	}

	funds := cfg.Funds
	currencies := cfg.Currencies

	var fundcodes []string
	for code := range funds {
		fundcodes = append(fundcodes, code)
	}

	fundresults, err := GetFunds(ctx, fundcodes...)
	if err != nil {
		if err == ErrDisabled {
			return nil
		}
		return err
	}

	c := notion.NewClient(secret)

	for _, r := range fundresults {
		pageID := funds[r.Code]

		if r.Price == 0 {
			continue
		}

		price := fmt.Sprintf("%.6f", r.Price)
		params := notion.UpdatePageParams{
			DatabasePageProperties: &notion.DatabasePageProperties{
				"Current Price ₺": notion.DatabasePageProperty{
					Type: notion.DBPropTypeRichText,
					RichText: []notion.RichText{
						{
							Type: notion.RichTextTypeText,
							Text: &notion.Text{Content: price},
						},
					},
				},
			},
		}
		_, err := c.UpdatePageProps(ctx, pageID, params)
		if err != nil {
			return err
		}
	}

	for code, pageID := range currencies {
		v, err := request(code)
		if err != nil {
			return err
		}

		price := fmt.Sprintf("%.6f", v)
		params := notion.UpdatePageParams{
			DatabasePageProperties: &notion.DatabasePageProperties{
				"Current Price ₺": notion.DatabasePageProperty{
					Type: notion.DBPropTypeRichText,
					RichText: []notion.RichText{
						{
							Type: notion.RichTextTypeText,
							Text: &notion.Text{Content: price},
						},
					},
				},
			},
		}
		_, err = c.UpdatePageProps(ctx, pageID, params)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>cmd/notion-yt: add logger<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/dstotijn/go-notion"
	"github.com/pelletier/go-toml"
)

func main() {
	if err := realmain(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

type Config struct {
	Currencies map[string]string
	Funds      map[string]string
}

func realmain() error {
	var (
		flagConfig = flag.String("c", "config.toml", "TOML configuration path")
	)
	flag.Parse()

	logger := log.New(os.Stdout, "", 0)

	bytes, err := ioutil.ReadFile(*flagConfig)
	if err != nil {
		return err
	}

	var cfg Config
	if err := toml.Unmarshal(bytes, &cfg); err != nil {
		return err
	}

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	secret := os.Getenv("NOTION_YT_SECRET")
	if secret == "" {
		return fmt.Errorf("NOTION_YET_SECRET environment variable is not set")
	}

	funds := cfg.Funds
	currencies := cfg.Currencies

	var fundcodes []string
	for code := range funds {
		fundcodes = append(fundcodes, code)
	}

	fundresults, err := GetFunds(ctx, fundcodes...)
	if err != nil {
		if err == ErrDisabled {
			return nil
		}
		return err
	}

	c := notion.NewClient(secret)

	for _, r := range fundresults {
		pageID := funds[r.Code]

		if r.Price == 0 {
			continue
		}

		price := fmt.Sprintf("%.6f", r.Price)
		params := notion.UpdatePageParams{
			DatabasePageProperties: &notion.DatabasePageProperties{
				"Current Price ₺": notion.DatabasePageProperty{
					Type: notion.DBPropTypeRichText,
					RichText: []notion.RichText{
						{
							Type: notion.RichTextTypeText,
							Text: &notion.Text{Content: price},
						},
					},
				},
			},
		}

		logger.Printf("Updating %q...", r.Code)
		_, err := c.UpdatePageProps(ctx, pageID, params)
		if err != nil {
			return err
		}
	}

	for code, pageID := range currencies {
		v, err := request(code)
		if err != nil {
			return err
		}

		price := fmt.Sprintf("%.6f", v)
		params := notion.UpdatePageParams{
			DatabasePageProperties: &notion.DatabasePageProperties{
				"Current Price ₺": notion.DatabasePageProperty{
					Type: notion.DBPropTypeRichText,
					RichText: []notion.RichText{
						{
							Type: notion.RichTextTypeText,
							Text: &notion.Text{Content: price},
						},
					},
				},
			},
		}

		logger.Printf("Updating %q...", code)
		_, err = c.UpdatePageProps(ctx, pageID, params)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package lv

import "sync"

// NewSpace returns an N-dimensional vector space.
func NewSpace() *Space {
	return &Space{}
}

// Space represents an N-dimensional vector space. Each name and unique label
// value pair establishes a new dimension and point within that dimension. Order
// matters, i.e. [a=1 b=2] identifies a different timeseries than [b=2 a=1].
type Space struct {
	mtx   sync.RWMutex
	nodes map[string]*node
}

// Observe locates the time series identified by the name and label values in
// the vector space, and appends the value to the list of observations.
func (s *Space) Observe(name string, lvs LabelValues, value float64) {
	s.nodeFor(name).observe(lvs, value)
}

// Add locates the time series identified by the name and label values in
// the vector space, and appends the delta to the last value in the list of
// observations.
func (s *Space) Add(name string, lvs LabelValues, delta float64) {
	s.nodeFor(name).add(lvs, delta)
}

// Walk traverses the vector space and invokes fn for each non-empty time series
// which is encountered. Return false to abort the traversal.
func (s *Space) Walk(fn func(name string, lvs LabelValues, observations []float64) bool) {
	s.mtx.RLock()
	defer s.mtx.RUnlock()
	for name, node := range s.nodes {
		f := func(lvs LabelValues, observations []float64) bool { return fn(name, lvs, observations) }
		if !node.walk(LabelValues{}, f) {
			return
		}
	}
}

// Reset empties the current space and returns a new Space with the old
// contents. Reset a Space to get an immutable copy suitable for walking.
func (s *Space) Reset() *Space {
	s.mtx.Lock()
	defer s.mtx.Unlock()
	n := NewSpace()
	n.nodes, s.nodes = s.nodes, n.nodes
	return n
}

func (s *Space) nodeFor(name string) *node {
	s.mtx.Lock()
	defer s.mtx.Unlock()
	if s.nodes == nil {
		s.nodes = map[string]*node{}
	}
	n, ok := s.nodes[name]
	if !ok {
		n = &node{}
		s.nodes[name] = n
	}
	return n
}

// node exists at a specific point in the N-dimensional vector space of all
// possible label values. The node collects observations and has child nodes
// with greater specificity.
type node struct {
	mtx          sync.RWMutex
	observations []float64
	children     map[pair]*node
}

type pair struct{ label, value string }

func (n *node) observe(lvs LabelValues, value float64) {
	n.mtx.Lock()
	defer n.mtx.Unlock()
	if len(lvs) == 0 {
		n.observations = append(n.observations, value)
		return
	}
	if len(lvs) < 2 {
		panic("too few LabelValues; programmer error!")
	}
	head, tail := pair{lvs[0], lvs[1]}, lvs[2:]
	if n.children == nil {
		n.children = map[pair]*node{}
	}
	child, ok := n.children[head]
	if !ok {
		child = &node{}
		n.children[head] = child
	}
	child.observe(tail, value)
}

func (n *node) add(lvs LabelValues, delta float64) {
	n.mtx.Lock()
	defer n.mtx.Unlock()
	if len(lvs) == 0 {
		var value float64
		if len(n.observations) > 0 {
			value = last(n.observations) + delta
		} else {
			value = delta
		}
		n.observations = append(n.observations, value)
		return
	}
	if len(lvs) < 2 {
		panic("too few LabelValues; programmer error!")
	}
	head, tail := pair{lvs[0], lvs[1]}, lvs[2:]
	if n.children == nil {
		n.children = map[pair]*node{}
	}
	child, ok := n.children[head]
	if !ok {
		child = &node{}
		n.children[head] = child
	}
	child.add(tail, delta)
}

func (n *node) walk(lvs LabelValues, fn func(LabelValues, []float64) bool) bool {
	n.mtx.RLock()
	defer n.mtx.RUnlock()
	if len(n.observations) > 0 && !fn(lvs, n.observations) {
		return false
	}
	for p, child := range n.children {
		if !child.walk(append(lvs, p.label, p.value), fn) {
			return false
		}
	}
	return true
}

func last(a []float64) float64 {
	return a[len(a)-1]
}
<commit_msg>metrics/internal/lv: fix len comparison<commit_after>package lv

import "sync"

// NewSpace returns an N-dimensional vector space.
func NewSpace() *Space {
	return &Space{}
}

// Space represents an N-dimensional vector space. Each name and unique label
// value pair establishes a new dimension and point within that dimension. Order
// matters, i.e. [a=1 b=2] identifies a different timeseries than [b=2 a=1].
type Space struct {
	mtx   sync.RWMutex
	nodes map[string]*node
}

// Observe locates the time series identified by the name and label values in
// the vector space, and appends the value to the list of observations.
func (s *Space) Observe(name string, lvs LabelValues, value float64) {
	s.nodeFor(name).observe(lvs, value)
}

// Add locates the time series identified by the name and label values in
// the vector space, and appends the delta to the last value in the list of
// observations.
func (s *Space) Add(name string, lvs LabelValues, delta float64) {
	s.nodeFor(name).add(lvs, delta)
}

// Walk traverses the vector space and invokes fn for each non-empty time series
// which is encountered. Return false to abort the traversal.
func (s *Space) Walk(fn func(name string, lvs LabelValues, observations []float64) bool) {
	s.mtx.RLock()
	defer s.mtx.RUnlock()
	for name, node := range s.nodes {
		f := func(lvs LabelValues, observations []float64) bool { return fn(name, lvs, observations) }
		if !node.walk(LabelValues{}, f) {
			return
		}
	}
}

// Reset empties the current space and returns a new Space with the old
// contents. Reset a Space to get an immutable copy suitable for walking.
func (s *Space) Reset() *Space {
	s.mtx.Lock()
	defer s.mtx.Unlock()
	n := NewSpace()
	n.nodes, s.nodes = s.nodes, n.nodes
	return n
}

func (s *Space) nodeFor(name string) *node {
	s.mtx.Lock()
	defer s.mtx.Unlock()
	if s.nodes == nil {
		s.nodes = map[string]*node{}
	}
	n, ok := s.nodes[name]
	if !ok {
		n = &node{}
		s.nodes[name] = n
	}
	return n
}

// node exists at a specific point in the N-dimensional vector space of all
// possible label values. The node collects observations and has child nodes
// with greater specificity.
type node struct {
	mtx          sync.RWMutex
	observations []float64
	children     map[pair]*node
}

type pair struct{ label, value string }

func (n *node) observe(lvs LabelValues, value float64) {
	n.mtx.Lock()
	defer n.mtx.Unlock()
	if len(lvs) <= 0 {
		n.observations = append(n.observations, value)
		return
	}
	if len(lvs) < 2 {
		panic("too few LabelValues; programmer error!")
	}
	head, tail := pair{lvs[0], lvs[1]}, lvs[2:]
	if n.children == nil {
		n.children = map[pair]*node{}
	}
	child, ok := n.children[head]
	if !ok {
		child = &node{}
		n.children[head] = child
	}
	child.observe(tail, value)
}

func (n *node) add(lvs LabelValues, delta float64) {
	n.mtx.Lock()
	defer n.mtx.Unlock()
	if len(lvs) <= 0 {
		var value float64
		if len(n.observations) > 0 {
			value = last(n.observations) + delta
		} else {
			value = delta
		}
		n.observations = append(n.observations, value)
		return
	}
	if len(lvs) < 2 {
		panic("too few LabelValues; programmer error!")
	}
	head, tail := pair{lvs[0], lvs[1]}, lvs[2:]
	if n.children == nil {
		n.children = map[pair]*node{}
	}
	child, ok := n.children[head]
	if !ok {
		child = &node{}
		n.children[head] = child
	}
	child.add(tail, delta)
}

func (n *node) walk(lvs LabelValues, fn func(LabelValues, []float64) bool) bool {
	n.mtx.RLock()
	defer n.mtx.RUnlock()
	if len(n.observations) > 0 && !fn(lvs, n.observations) {
		return false
	}
	for p, child := range n.children {
		if !child.walk(append(lvs, p.label, p.value), fn) {
			return false
		}
	}
	return true
}

func last(a []float64) float64 {
	return a[len(a)-1]
}
<|endoftext|>


<commit_before>// +build linux

package sysstats

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// CpuRawStats represents *one* CPU raw statistics of a linux system.
type CpuRawStats map[string]uint64

// CpuStats represents *one* CPU statistics of a linux system.
type CpuStats map[string]float64

// CpusRawStats represents *all* the CPU raw statistics of a linux system.
type CpusRawStats map[string]CpuRawStats

// CpusStats represents *all* the CPU statistics of a linux system.
type CpusStats map[string]CpuStats

// getCpuRawStats gets the Cpu raw stats of a linux system from the
// file /proc/stat
// It retuns a CpusRawStats var. CpusRawStats is a map which keys are the
// 'kernel name' of the CPUs (cpu, cpu0, etc). The value of each key is a
// CpuRawStats var with the statistics for that CPU.
// CpusRawStats has the following key:
//  Name      - Name of the Cpu (as it is on /proc/stat: cpu, cpu0,...).
// CpuRawStats has the following keys:
//  User      - % Cpu time spent in user mode.
//  Nice      - % Cpu time spent in user mode with low priority (nice).
//  System    - % Cpu time spent in system mode.
//  Idle      - % Cpu time spent in the idle task.
//  Iowait    - % Cpu time spent waiting for I/O to complete (since 2.5.41).
//  Irq       - % Cpu time servicing interrupts (since 2.6.0-test4).
//  Softirq   - % Cpu time servicing softirqs (since 2.6.0-test4).
//  Steal     - % Cpu stolen time, which is the time spent in other operating
//              systems when running a virtualized environment (since 2.6.11).
//  Guest     - % Cpu time spent running a virtual Cpu for guest operating
//              systems under the control of the Linux kernel (since 2.6.24).
//  GuestNice - % Cpu time spent running a niced guest (virtual Cpu for guest
//              operating systems under the control of the Linux kernel)
//              (since 2.6.33).
//  Total     - % Cpu utilization (not idle).
func getCpuRawStats() (cpusRawStats CpusRawStats, err error) {
	file, err := os.Open("/proc/stat")
	if err != nil {
		return nil, err
	}
	defer file.Close()

	cpusRawStats = CpusRawStats{}

	re := regexp.MustCompile(`^cpu.*$`)

	scanner := bufio.NewScanner(file)
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		stats := re.FindString(line)
		if stats == "" {
			// No match so no more cpu 'lines'
			break
		}
		fmt.Println(stats)
		cpuName, rawStats, err := parseCpuRawStats(stats)
		if err != nil {
			return nil, err
		}
		cpusRawStats[cpuName] = rawStats
	}

	return cpusRawStats, nil
}

// parseCpuRawStats parses the CPU stats as they are in the file /proc/stat
// stats has the following format:
//   cpu  294 0 309 10612 71 30 0 0 0 0
// cpuName is the name of the CPU (cpu, cpu0, cpu1, etc) and rawStats has
// the following format:
//   map[User:9366 Nice:0 System:5692 Iowait:114 Steal:0 GuestNice:0
//       Idle:1458880 Irq:806 Softirq:0 Guest:0]
func parseCpuRawStats(stats string) (cpuName string, rawStats CpuRawStats,
	err error) {
	rawStats = CpuRawStats{}

	fields := strings.Fields(stats)
	cpuName = fields[0]
	for i := 1; i < len(fields); i++ {
		stat, err := strconv.ParseUint(fields[i], 10, 64)
		if err != nil {
			return "", nil, err
		}
		switch i {
		case 1:
			rawStats[`User`] = stat
		case 2:
			rawStats[`Nice`] = stat
		case 3:
			rawStats[`System`] = stat
		case 4:
			rawStats[`Idle`] = stat
		case 5:
			rawStats[`Iowait`] = stat
		case 6:
			rawStats[`Irq`] = stat
		case 7:
			rawStats[`Softirq`] = stat
		case 8:
			rawStats[`Steal`] = stat
		case 9:
			rawStats[`Guest`] = stat
		case 10:
			rawStats[`GuestNice`] = stat
		}
	}

	return cpuName, rawStats, nil
}

func getCpuStats() (cpuStats CpusRawStats, err error) {
	return getCpuRawStats()
}
<commit_msg>Raw statistics measuresw time instead of the % of CPU utilization<commit_after>// +build linux

package sysstats

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// CpuRawStats represents *one* CPU raw statistics of a linux system.
type CpuRawStats map[string]uint64

// CpuStats represents *one* CPU statistics of a linux system.
type CpuStats map[string]float64

// CpusRawStats represents *all* the CPU raw statistics of a linux system.
type CpusRawStats map[string]CpuRawStats

// CpusStats represents *all* the CPU statistics of a linux system.
type CpusStats map[string]CpuStats

// getCpuRawStats gets the Cpu raw stats of a linux system from the
// file /proc/stat
// It retuns a CpusRawStats var. CpusRawStats is a map which keys are the
// 'kernel name' of the CPUs (cpu, cpu0, etc). The value of each key is a
// CpuRawStats var with the statistics for that CPU.
// CpusRawStats has the following key:
//  Name      - Name of the Cpu (as it is on /proc/stat: cpu, cpu0,...).
// CpuRawStats has the following keys:
//  User      - Time spent in user mode.
//  Nice      - Time spent in user mode with low priority (nice).
//  System    - Time spent in system mode.
//  Idle      - Time spent in the idle task.
//  Iowait    - Time spent waiting for I/O to complete (since 2.5.41).
//  Irq       - Time servicing interrupts (since 2.6.0-test4).
//  Softirq   - Time servicing softirqs (since 2.6.0-test4).
//  Steal     - Stolen time, which is the time spent in other operating
//              systems when running a virtualized environment (since 2.6.11).
//  Guest     - Time spent running a virtual Cpu for guest operating
//              systems under the control of the Linux kernel (since 2.6.24).
//  GuestNice - Time spent running a niced guest (virtual Cpu for guest
//              operating systems under the control of the Linux kernel)
//              (since 2.6.33).
//  Total     - Total time (not idle).
func getCpuRawStats() (cpusRawStats CpusRawStats, err error) {
	file, err := os.Open("/proc/stat")
	if err != nil {
		return nil, err
	}
	defer file.Close()

	cpusRawStats = CpusRawStats{}

	re := regexp.MustCompile(`^cpu.*$`)

	scanner := bufio.NewScanner(file)
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		stats := re.FindString(line)
		if stats == "" {
			// No match so no more cpu 'lines'
			break
		}
		fmt.Println(stats)
		cpuName, rawStats, err := parseCpuRawStats(stats)
		if err != nil {
			return nil, err
		}
		cpusRawStats[cpuName] = rawStats
	}

	return cpusRawStats, nil
}

// parseCpuRawStats parses the CPU stats as they are in the file /proc/stat
// stats has the following format:
//   cpu  294 0 309 10612 71 30 0 0 0 0
// cpuName is the name of the CPU (cpu, cpu0, cpu1, etc) and rawStats has
// the following format:
//   map[User:9366 Nice:0 System:5692 Iowait:114 Steal:0 GuestNice:0
//       Idle:1458880 Irq:806 Softirq:0 Guest:0]
func parseCpuRawStats(stats string) (cpuName string, rawStats CpuRawStats,
	err error) {
	rawStats = CpuRawStats{}

	fields := strings.Fields(stats)
	cpuName = fields[0]
	for i := 1; i < len(fields); i++ {
		stat, err := strconv.ParseUint(fields[i], 10, 64)
		if err != nil {
			return "", nil, err
		}
		switch i {
		case 1:
			rawStats[`User`] = stat
		case 2:
			rawStats[`Nice`] = stat
		case 3:
			rawStats[`System`] = stat
		case 4:
			rawStats[`Idle`] = stat
		case 5:
			rawStats[`Iowait`] = stat
		case 6:
			rawStats[`Irq`] = stat
		case 7:
			rawStats[`Softirq`] = stat
		case 8:
			rawStats[`Steal`] = stat
		case 9:
			rawStats[`Guest`] = stat
		case 10:
			rawStats[`GuestNice`] = stat
		}
	}

	return cpuName, rawStats, nil
}

func getCpuStats() (cpuStats CpusRawStats, err error) {
	return getCpuRawStats()
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "golang.org/x/net/html"
    "net/http"
    "os"
    "strings"
)

// Helper function to pull the href attribute from a Token
func getHref(t html.Token) (ok bool, href string) {
    // Iterate over all of the Token's attributes until we find an "href"
    for _, a := range t.Attr {
        if a.Key == "href" {
            href = a.Val
            ok = true
        }
    }

    // "bare" return will return the variables (ok, href) as defined in
    // the function definition
    return
}

// Extract all http** links from a given webpage
func crawl(url string, ch chan string, chFinished chan bool) {
    resp, err := http.Get(url)

    defer func() {
        // Notify that we're done after this function
        chFinished <- true
    }()

    if err != nil {
        fmt.Println("ERROR: Failed to crawl \"" + url + "\"")
        return
    }

    b := resp.Body
    defer b.Close() // close Body when the function returns

    z := html.NewTokenizer(b)

    for {
        tt := z.Next()

        switch {
        case tt == html.ErrorToken:
            // End of the document, we're done
            return
        case tt == html.StartTagToken:
            t := z.Token()

            // Check if the token is an <a> tag
            isAnchor := t.Data == "a"
            if !isAnchor {
                continue
            }

            // Extract the href value, if there is one
            ok, url := getHref(t)
            if !ok {
                continue
            }

            // Make sure the url begines in http**
            hasProto := strings.Index(url, "http") == 0
            if hasProto {
                ch <- url
            }
        }
    }
}

func main() {
    foundUrls := make(map[string]bool)
    seedUrls := os.Args[1:]

    // Channels
    chUrls := make(chan string)
    chFinished := make(chan bool)

    // Kick off the crawl process (concurrently)
    for _, url := range seedUrls {
        go crawl(url, chUrls, chFinished)
    }

    // Subscribe to both channels
    for c := 0; c < len(seedUrls); {
        select {
        case url := <-chUrls:
            foundUrls[url] = true
        case <-chFinished:
            c++
        }
    }

    // We're done! Print the results...

    fmt.Println("\nFound", len(foundUrls), "unique urls:\n")

    for url, _ := range foundUrls {
        fmt.Println(" - " + url)
    }

    close(chUrls)
}<commit_msg>changed function names<commit_after>package main

import (
    "fmt"
    "golang.org/x/net/html"
    "net/http"
    "os"
    "strings"
)

// Helper function to pull the href attribute from a Token
func getHref(t html.Token) (ok bool, href string) {
    // Iterate over all of the Token's attributes until we find an "href"
    for _, a := range t.Attr {
        if a.Key == "href" {
            href = a.Val
            ok = true
        }
    }

    // "bare" return will return the variables (ok, href) as defined in
    // the function definition
    return
}

// Extract all http** links from a given webpage
func scrape(url string, ch chan string, chFinished chan bool) {
    resp, err := http.Get(url)

    defer func() {
        // Notify that we're done after this function
        chFinished <- true
    }()

    if err != nil {
        fmt.Println("ERROR: Failed to crawl \"" + url + "\"")
        return
    }

    b := resp.Body
    defer b.Close() // close Body when the function returns

    z := html.NewTokenizer(b)

    for {
        tt := z.Next()

        switch {
        case tt == html.ErrorToken:
            // End of the document, we're done
            return
        case tt == html.StartTagToken:
            t := z.Token()

            // Check if the token is an <a> tag
            isAnchor := t.Data == "a"
            if !isAnchor {
                continue
            }

            // Extract the href value, if there is one
            ok, url := getHref(t)
            if !ok {
                continue
            }

            // Make sure the url begines in http**
            hasProto := strings.Index(url, "http") == 0
            if hasProto {
                ch <- url
            }
        }
    }
}

func main() {
    foundUrls := make(map[string]bool)
    seedUrls := os.Args[1:]

    // Channels
    chUrls := make(chan string)
    chFinished := make(chan bool)

    // Kick off the crawl process (concurrently)
    for _, url := range seedUrls {
        go scrape(url, chUrls, chFinished)
    }

    // Subscribe to both channels
    for c := 0; c < len(seedUrls); {
        select {
        case url := <-chUrls:
            foundUrls[url] = true
        case <-chFinished:
            c++
        }
    }

    // We're done! Print the results...

    fmt.Println("\nFound", len(foundUrls), "unique urls:\n")

    for url, _ := range foundUrls {
        fmt.Println(" - " + url)
    }

    close(chUrls)
}<|endoftext|>


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

// +build !windows

// Test that the Go runtime still works if C code changes the signal stack.

package cgotest

/*
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static stack_t oss;
static char signalStack[SIGSTKSZ];

static void changeSignalStack() {
	stack_t ss;
	memset(&ss, 0, sizeof ss);
	ss.ss_sp = signalStack;
	ss.ss_flags = 0;
	ss.ss_size = SIGSTKSZ;
	if (sigaltstack(&ss, &oss) < 0) {
		perror("sigaltstack");
		abort();
	}
}

static void restoreSignalStack() {
#if (defined(__x86_64__) || defined(__i386__)) && defined(__APPLE__)
	// The Darwin C library enforces a minimum that the kernel does not.
	// This is OK since we allocated this much space in mpreinit,
	// it was just removed from the buffer by stackalloc.
	oss.ss_size = MINSIGSTKSZ;
#endif
	if (sigaltstack(&oss, NULL) < 0) {
		perror("sigaltstack restore");
		abort();
	}
}

static int zero() {
	return 0;
}
*/
import "C"

import (
	"runtime"
	"testing"
)

func testSigaltstack(t *testing.T) {
	switch {
	case runtime.GOOS == "solaris", runtime.GOOS == "darwin" && (runtime.GOARCH == "arm" || runtime.GOARCH == "arm64"):
		t.Skipf("switching signal stack not implemented on %s/%s", runtime.GOOS, runtime.GOARCH)
	}

	C.changeSignalStack()
	defer C.restoreSignalStack()
	defer func() {
		if recover() == nil {
			t.Error("did not see expected panic")
		}
	}()
	v := 1 / int(C.zero())
	t.Errorf("unexpected success of division by zero == %d", v)
}
<commit_msg>misc/cgo/test: disable sigaltstack test on darwin/386<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !windows

// Test that the Go runtime still works if C code changes the signal stack.

package cgotest

/*
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static stack_t oss;
static char signalStack[SIGSTKSZ];

static void changeSignalStack() {
	stack_t ss;
	memset(&ss, 0, sizeof ss);
	ss.ss_sp = signalStack;
	ss.ss_flags = 0;
	ss.ss_size = SIGSTKSZ;
	if (sigaltstack(&ss, &oss) < 0) {
		perror("sigaltstack");
		abort();
	}
}

static void restoreSignalStack() {
#if (defined(__x86_64__) || defined(__i386__)) && defined(__APPLE__)
	// The Darwin C library enforces a minimum that the kernel does not.
	// This is OK since we allocated this much space in mpreinit,
	// it was just removed from the buffer by stackalloc.
	oss.ss_size = MINSIGSTKSZ;
#endif
	if (sigaltstack(&oss, NULL) < 0) {
		perror("sigaltstack restore");
		abort();
	}
}

static int zero() {
	return 0;
}
*/
import "C"

import (
	"runtime"
	"testing"
)

func testSigaltstack(t *testing.T) {
	switch {
	case runtime.GOOS == "solaris", runtime.GOOS == "darwin" && (runtime.GOARCH == "arm" || runtime.GOARCH == "arm64"):
		t.Skipf("switching signal stack not implemented on %s/%s", runtime.GOOS, runtime.GOARCH)
	case runtime.GOOS == "darwin" && runtime.GOARCH == "386":
		t.Skipf("sigaltstack fails on darwin/386")
	}

	C.changeSignalStack()
	defer C.restoreSignalStack()
	defer func() {
		if recover() == nil {
			t.Error("did not see expected panic")
		}
	}()
	v := 1 / int(C.zero())
	t.Errorf("unexpected success of division by zero == %d", v)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/binary"
	"hash"
	"hash/crc64"
	"io"
	"os"
	"path/filepath"
	"sync"
	"syscall"
)

func directoryScanner(directoryScanQueue chan string, fileReadQueue chan string, blockQueue chan block, excludePatterns []string, workInProgress *sync.WaitGroup) {
	for directoryPath := range directoryScanQueue {
		if verbose {
			logger.Println(directoryPath)
		}

		directory, err := os.Open(directoryPath)
		if err != nil {
			logger.Println("directory read error:", err.Error())
			workInProgress.Done()
			continue
		}

		uid, gid, mode := getModeOwnership(directory)
		blockQueue <- block{directoryPath, 0, nil, blockTypeDirectory, uid, gid, mode}

		for fileName := range readdirnames(int(directory.Fd())) {
			filePath := filepath.Join(directoryPath, fileName)

			excludeFile := false
			for _, excludePattern := range excludePatterns {
				match, err := filepath.Match(excludePattern, filePath)
				if err == nil && match {
					excludeFile = true
					break
				}
			}
			if excludeFile {
				logger.Println("skipping excluded file", filePath)
				continue
			}

			fileInfo, err := os.Lstat(filePath)
			if err != nil {
				logger.Println("unable to lstat file", err.Error())
				continue
			} else if (fileInfo.Mode() & os.ModeSymlink) != 0 {
				logger.Println("skipping symbolic link", filePath)
				continue
			}

			workInProgress.Add(1)
			if fileInfo.IsDir() {
				directoryScanQueue <- filePath
			} else {
				fileReadQueue <- filePath
			}
		}

		directory.Close()
		workInProgress.Done()
	}
}

func getModeOwnership(file *os.File) (int, int, os.FileMode) {
	var uid int = 0
	var gid int = 0
	var mode os.FileMode = 0
	fi, err := file.Stat()
	if err != nil {
		logger.Println("file stat error; uid/gid/mode will be incorrect:", err.Error())
	} else {
		mode = fi.Mode()
		stat_t := fi.Sys().(*syscall.Stat_t)
		if stat_t != nil {
			uid = int(stat_t.Uid)
			gid = int(stat_t.Gid)
		} else {
			logger.Println("unable to find file uid/gid")
		}
	}
	return uid, gid, mode
}

func fileReader(fileReadQueue <-chan string, blockQueue chan block, workInProgress *sync.WaitGroup) {
	for filePath := range fileReadQueue {
		if verbose {
			logger.Println(filePath)
		}

		file, err := os.Open(filePath)
		if err == nil {

			uid, gid, mode := getModeOwnership(file)
			blockQueue <- block{filePath, 0, nil, blockTypeStartOfFile, uid, gid, mode}

			bufferedFile := bufio.NewReader(file)

			for {
				buffer := make([]byte, blockSize)
				bytesRead, err := bufferedFile.Read(buffer)
				if err == io.EOF {
					break
				} else if err != nil {
					logger.Println("file read error; file contents will be incomplete:", err.Error())
					break
				}

				blockQueue <- block{filePath, uint16(bytesRead), buffer, blockTypeData, 0, 0, 0}
			}

			blockQueue <- block{filePath, 0, nil, blockTypeEndOfFile, 0, 0, 0}
			file.Close()
		} else {
			logger.Println("file open error:", err.Error())
		}

		workInProgress.Done()
	}
}

func archiveWriter(output io.Writer, blockQueue <-chan block) {
	hash := crc64.New(crc64.MakeTable(crc64.ECMA))
	output = io.MultiWriter(output, hash)
	blockCount := 0
	blockType := make([]byte, 1)

	// Archive header: stole ideas from the PNG file header here, but replaced
	// 'PNG' with 'FA1' to identify the fast-archive version 1 format.
	_, err := output.Write([]byte{0x89, 0x46, 0x41, 0x31, 0x0D, 0x0A, 0x1A, 0x0A})
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}

	for block := range blockQueue {
		filePath := []byte(block.filePath)
		err = binary.Write(output, binary.BigEndian, uint16(len(filePath)))
		if err != nil {
			logger.Fatalln("Archive write error:", err.Error())
		}
		_, err = output.Write(filePath)
		if err != nil {
			logger.Fatalln("Archive write error:", err.Error())
		}

		if block.blockType == blockTypeStartOfFile {
			blockType[0] = byte(blockTypeStartOfFile)
			_, err = output.Write(blockType)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, uint32(block.uid))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, uint32(block.gid))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, block.mode)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
		} else if block.blockType == blockTypeEndOfFile {
			blockType[0] = byte(blockTypeEndOfFile)
			_, err = output.Write(blockType)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
		} else if block.blockType == blockTypeData {
			blockType[0] = byte(blockTypeData)
			_, err = output.Write(blockType)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}

			err = binary.Write(output, binary.BigEndian, uint16(block.numBytes))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}

			_, err = output.Write(block.buffer[:block.numBytes])
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
		} else if block.blockType == blockTypeDirectory {
			blockType[0] = byte(blockTypeDirectory)
			_, err = output.Write(blockType)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, uint32(block.uid))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, uint32(block.gid))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, block.mode)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
		} else {
			logger.Panicln("Unexpected block type")
		}

		blockCount += 1
		if (blockCount % 1000) == 0 {
			writeChecksumBlock(hash, output, blockType)
		}
	}

	writeChecksumBlock(hash, output, blockType)
}

func writeChecksumBlock(hash hash.Hash64, output io.Writer, blockType []byte) {
	// file path length... zero
	err := binary.Write(output, binary.BigEndian, uint16(0))
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}

	blockType[0] = byte(blockTypeChecksum)
	_, err = output.Write(blockType)
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}
	binary.Write(output, binary.BigEndian, hash.Sum64())
}

// Copy of os.Readdirnames for UNIX systems, but modified to return results
// as found through a channel rather than in one large array.
func readdirnames(fd int) chan string {
	retval := make(chan string)
	go func(fd int) {
		var buf []byte = make([]byte, blockSize)
		var nbuf int
		var bufp int

		for {
			// Refill the buffer if necessary
			if bufp >= nbuf {
				bufp = 0
				var errno error
				nbuf, errno = syscall.ReadDirent(fd, buf)
				if errno != nil {
					err := os.NewSyscallError("readdirent", errno)
					logger.Println("error reading directory:", err.Error())
					break
				}
				if nbuf <= 0 {
					break // EOF
				}
			}

			// Drain the buffer
			var nb, nc int
			names := make([]string, 0, 100)
			nb, nc, names = syscall.ParseDirent(buf[bufp:nbuf], -1, names)
			bufp += nb

			for i := 0; i < nc; i++ {
				retval <- names[i]
			}
		}

		close(retval)
	}(fd)
	return retval
}
<commit_msg>Fix Issue #4; prevent deadlock in directoryScanner<commit_after>package main

import (
	"bufio"
	"encoding/binary"
	"hash"
	"hash/crc64"
	"io"
	"os"
	"path/filepath"
	"sync"
	"syscall"
)

func directoryScanner(directoryScanQueue chan string, fileReadQueue chan string, blockQueue chan block, excludePatterns []string, workInProgress *sync.WaitGroup) {
	for directoryPath := range directoryScanQueue {
		if verbose {
			logger.Println(directoryPath)
		}

		directory, err := os.Open(directoryPath)
		if err != nil {
			logger.Println("directory read error:", err.Error())
			workInProgress.Done()
			continue
		}

		uid, gid, mode := getModeOwnership(directory)
		blockQueue <- block{directoryPath, 0, nil, blockTypeDirectory, uid, gid, mode}

		for fileName := range readdirnames(int(directory.Fd())) {
			filePath := filepath.Join(directoryPath, fileName)

			excludeFile := false
			for _, excludePattern := range excludePatterns {
				match, err := filepath.Match(excludePattern, filePath)
				if err == nil && match {
					excludeFile = true
					break
				}
			}
			if excludeFile {
				logger.Println("skipping excluded file", filePath)
				continue
			}

			fileInfo, err := os.Lstat(filePath)
			if err != nil {
				logger.Println("unable to lstat file", err.Error())
				continue
			} else if (fileInfo.Mode() & os.ModeSymlink) != 0 {
				logger.Println("skipping symbolic link", filePath)
				continue
			}

			workInProgress.Add(1)
			if fileInfo.IsDir() {
				// Sending to directoryScanQueue can block if it's full; since
				// we're also the goroutine responsible for reading from it,
				// this could cause a deadlock.  We break that deadlock by
				// performing the send in a goroutine, where it can block
				// safely.  This does have the side-effect that
				// directoryScanQueue's max size is pretty much ineffective...
				// but that's better than a deadlock.
				go func(filePath string) {
					directoryScanQueue <- filePath
				}(filePath)
			} else {
				fileReadQueue <- filePath
			}
		}

		directory.Close()
		workInProgress.Done()
	}
}

func getModeOwnership(file *os.File) (int, int, os.FileMode) {
	var uid int = 0
	var gid int = 0
	var mode os.FileMode = 0
	fi, err := file.Stat()
	if err != nil {
		logger.Println("file stat error; uid/gid/mode will be incorrect:", err.Error())
	} else {
		mode = fi.Mode()
		stat_t := fi.Sys().(*syscall.Stat_t)
		if stat_t != nil {
			uid = int(stat_t.Uid)
			gid = int(stat_t.Gid)
		} else {
			logger.Println("unable to find file uid/gid")
		}
	}
	return uid, gid, mode
}

func fileReader(fileReadQueue <-chan string, blockQueue chan block, workInProgress *sync.WaitGroup) {
	for filePath := range fileReadQueue {
		if verbose {
			logger.Println(filePath)
		}

		file, err := os.Open(filePath)
		if err == nil {

			uid, gid, mode := getModeOwnership(file)
			blockQueue <- block{filePath, 0, nil, blockTypeStartOfFile, uid, gid, mode}

			bufferedFile := bufio.NewReader(file)

			for {
				buffer := make([]byte, blockSize)
				bytesRead, err := bufferedFile.Read(buffer)
				if err == io.EOF {
					break
				} else if err != nil {
					logger.Println("file read error; file contents will be incomplete:", err.Error())
					break
				}

				blockQueue <- block{filePath, uint16(bytesRead), buffer, blockTypeData, 0, 0, 0}
			}

			blockQueue <- block{filePath, 0, nil, blockTypeEndOfFile, 0, 0, 0}
			file.Close()
		} else {
			logger.Println("file open error:", err.Error())
		}

		workInProgress.Done()
	}
}

func archiveWriter(output io.Writer, blockQueue <-chan block) {
	hash := crc64.New(crc64.MakeTable(crc64.ECMA))
	output = io.MultiWriter(output, hash)
	blockCount := 0
	blockType := make([]byte, 1)

	// Archive header: stole ideas from the PNG file header here, but replaced
	// 'PNG' with 'FA1' to identify the fast-archive version 1 format.
	_, err := output.Write([]byte{0x89, 0x46, 0x41, 0x31, 0x0D, 0x0A, 0x1A, 0x0A})
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}

	for block := range blockQueue {
		filePath := []byte(block.filePath)
		err = binary.Write(output, binary.BigEndian, uint16(len(filePath)))
		if err != nil {
			logger.Fatalln("Archive write error:", err.Error())
		}
		_, err = output.Write(filePath)
		if err != nil {
			logger.Fatalln("Archive write error:", err.Error())
		}

		if block.blockType == blockTypeStartOfFile {
			blockType[0] = byte(blockTypeStartOfFile)
			_, err = output.Write(blockType)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, uint32(block.uid))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, uint32(block.gid))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, block.mode)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
		} else if block.blockType == blockTypeEndOfFile {
			blockType[0] = byte(blockTypeEndOfFile)
			_, err = output.Write(blockType)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
		} else if block.blockType == blockTypeData {
			blockType[0] = byte(blockTypeData)
			_, err = output.Write(blockType)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}

			err = binary.Write(output, binary.BigEndian, uint16(block.numBytes))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}

			_, err = output.Write(block.buffer[:block.numBytes])
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
		} else if block.blockType == blockTypeDirectory {
			blockType[0] = byte(blockTypeDirectory)
			_, err = output.Write(blockType)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, uint32(block.uid))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, uint32(block.gid))
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
			err = binary.Write(output, binary.BigEndian, block.mode)
			if err != nil {
				logger.Fatalln("Archive write error:", err.Error())
			}
		} else {
			logger.Panicln("Unexpected block type")
		}

		blockCount += 1
		if (blockCount % 1000) == 0 {
			writeChecksumBlock(hash, output, blockType)
		}
	}

	writeChecksumBlock(hash, output, blockType)
}

func writeChecksumBlock(hash hash.Hash64, output io.Writer, blockType []byte) {
	// file path length... zero
	err := binary.Write(output, binary.BigEndian, uint16(0))
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}

	blockType[0] = byte(blockTypeChecksum)
	_, err = output.Write(blockType)
	if err != nil {
		logger.Fatalln("Archive write error:", err.Error())
	}
	binary.Write(output, binary.BigEndian, hash.Sum64())
}

// Copy of os.Readdirnames for UNIX systems, but modified to return results
// as found through a channel rather than in one large array.
func readdirnames(fd int) chan string {
	retval := make(chan string)
	go func(fd int) {
		var buf []byte = make([]byte, blockSize)
		var nbuf int
		var bufp int

		for {
			// Refill the buffer if necessary
			if bufp >= nbuf {
				bufp = 0
				var errno error
				nbuf, errno = syscall.ReadDirent(fd, buf)
				if errno != nil {
					err := os.NewSyscallError("readdirent", errno)
					logger.Println("error reading directory:", err.Error())
					break
				}
				if nbuf <= 0 {
					break // EOF
				}
			}

			// Drain the buffer
			var nb, nc int
			names := make([]string, 0, 100)
			nb, nc, names = syscall.ParseDirent(buf[bufp:nbuf], -1, names)
			bufp += nb

			for i := 0; i < nc; i++ {
				retval <- names[i]
			}
		}

		close(retval)
	}(fd)
	return retval
}
<|endoftext|>