<commit_before>package main

import (
    "fmt"
    "bytes"
    "strconv"
    "net/http"
    "io/ioutil"
    "encoding/json"
    "github.com/andlabs/ui"
)

func main() {
    err := ui.Main(func() {
        vBox := ui.NewVerticalBox()
        hBox := ui.NewHorizontalBox()

        urlEntry := ui.NewEntry()
        jsonLabel := ui.NewLabel("Input json:")
        jsonEntry := ui.NewEntry()
        goButton := ui.NewButton("Go")
        resLabel := ui.NewLabel("")
        selectCombobox := ui.NewCombobox()

        hBox.Append(selectCombobox, false)
        hBox.Append(urlEntry, true)
        hBox.Append(goButton, false)
        hBox.SetPadded(true)

        vBox.Append(hBox, false)

        vBox.Append(jsonLabel, false)
        vBox.Append(jsonEntry, false)
        vBox.Append(resLabel, false)

        vBox.SetPadded(true)

        jsonLabel.Hide()
        jsonEntry.Hide()

        selectCombobox.Append("Get")
        selectCombobox.Append("Post")

        selectCombobox.SetSelected(0)

        window := ui.NewWindow("doreq - do request", 800, 600, false)
        window.SetChild(vBox)
        window.SetMargined(true)

        selectCombobox.OnSelected(func(*ui.Combobox) {
            selectedIndex := selectCombobox.Selected()
            fmt.Println("selectedIndex:", selectedIndex)
            switch selectedIndex {
                case 1:
                    jsonLabel.Show()
                    jsonEntry.Show()
                default:
                    jsonLabel.Hide()
                    jsonEntry.Hide()
            }
        })

        goButton.OnClicked(func(*ui.Button) {
            fmt.Printf("goButton pressed\n")
            selectedIndex := selectCombobox.Selected()
            switch selectedIndex {
                case 0:
                    resp, err := http.Get(urlEntry.Text())
                    if err != nil {
                        resLabel.SetText(err.Error())
                        return
                    }
                    handleResp(resp, resLabel)
                case 1:
                    url := urlEntry.Text()
                    fmt.Println("URL:", url)
                    jsonStr := jsonEntry.Text()

                    var jsonBytes = []byte(jsonStr)
                    req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBytes))
                    req.Header.Set("Content-Type", "application/json")

                    client := &http.Client{}
                    resp, err := client.Do(req)
                    if err != nil {
                        panic(err)
                    }
                    handleResp(resp, resLabel)
                default:
                    fmt.Printf("no selected")
            }

        })
        window.OnClosing(func(*ui.Window) bool {
            ui.Quit()
            return true
        })
        window.Show()
    })
    if err != nil {
        panic(err)
    }
}

func isJSON(s string) bool {
    var js map[string]interface{}
    return json.Unmarshal([]byte(s), &js) == nil
}

func handleResp(resp *http.Response, label *ui.Label) {
    defer resp.Body.Close()
    fmt.Println("response Status:", resp.Status)
    fmt.Println("response Headers:", resp.Header)

    if resp.StatusCode == 200 {
        bodyBytes, err := ioutil.ReadAll(resp.Body)
        if err != nil {
            return
        }
        bodyString := string(bodyBytes)
        if !isJSON(bodyString) {
            bodyString = "Response is not json."
        }
        label.SetText(bodyString)
    } else {
        label.SetText("Error with code " + strconv.Itoa(resp.StatusCode))
    }
}<commit_msg>format<commit_after>package main

import (
    "fmt"
    "bytes"
    "strconv"
    "net/http"
    "io/ioutil"
    "encoding/json"
    "github.com/andlabs/ui"
)

func main() {
    err := ui.Main(func() {
        vBox := ui.NewVerticalBox()
        hBox := ui.NewHorizontalBox()

        urlEntry := ui.NewEntry()
        jsonLabel := ui.NewLabel("Input json:")
        jsonEntry := ui.NewEntry()
        goButton := ui.NewButton("Go")
        resLabel := ui.NewLabel("")
        selectCombobox := ui.NewCombobox()

        hBox.Append(selectCombobox, false)
        hBox.Append(urlEntry, true)
        hBox.Append(goButton, false)
        hBox.SetPadded(true)

        vBox.Append(hBox, false)
        vBox.Append(jsonLabel, false)
        vBox.Append(jsonEntry, false)
        vBox.Append(resLabel, false)
        vBox.SetPadded(true)

        jsonLabel.Hide()
        jsonEntry.Hide()

        selectCombobox.Append("Get")
        selectCombobox.Append("Post")
        selectCombobox.SetSelected(0)

        window := ui.NewWindow("doreq - do request", 800, 600, false)
        window.SetChild(vBox)
        window.SetMargined(true)

        selectCombobox.OnSelected(func(*ui.Combobox) {
            selectedIndex := selectCombobox.Selected()
            fmt.Println("selectedIndex:", selectedIndex)
            switch selectedIndex {
                case 1:
                    jsonLabel.Show()
                    jsonEntry.Show()
                default:
                    jsonLabel.Hide()
                    jsonEntry.Hide()
            }
        })

        goButton.OnClicked(func(*ui.Button) {
            fmt.Printf("goButton pressed\n")
            selectedIndex := selectCombobox.Selected()
            switch selectedIndex {
                case 0:
                    resp, err := http.Get(urlEntry.Text())
                    if err != nil {
                        resLabel.SetText(err.Error())
                        return
                    }
                    handleResp(resp, resLabel)
                case 1:
                    url := urlEntry.Text()
                    fmt.Println("URL:", url)
                    jsonStr := jsonEntry.Text()

                    var jsonBytes = []byte(jsonStr)
                    req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBytes))
                    req.Header.Set("Content-Type", "application/json")

                    client := &http.Client{}
                    resp, err := client.Do(req)
                    if err != nil {
                        panic(err)
                    }
                    handleResp(resp, resLabel)
                default:
                    fmt.Printf("no selected")
            }

        })
        window.OnClosing(func(*ui.Window) bool {
            ui.Quit()
            return true
        })
        window.Show()
    })
    if err != nil {
        panic(err)
    }
}

func isJSON(s string) bool {
    var js map[string]interface{}
    return json.Unmarshal([]byte(s), &js) == nil
}

func handleResp(resp *http.Response, label *ui.Label) {
    defer resp.Body.Close()
    fmt.Println("response Status:", resp.Status)
    fmt.Println("response Headers:", resp.Header)

    if resp.StatusCode == 200 {
        bodyBytes, err := ioutil.ReadAll(resp.Body)
        if err != nil {
            return
        }
        bodyString := string(bodyBytes)
        if !isJSON(bodyString) {
            bodyString = "Response is not json."
        }
        label.SetText(bodyString)
    } else {
        label.SetText("Error with code " + strconv.Itoa(resp.StatusCode))
    }
}<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"log"
	"net/http"
	"time"

	"github.com/bmizerany/lpx"
	"github.com/heroku/slog"
	"github.com/kr/logfmt"
)

var (
	TokenPrefix = []byte("t.")
	Heroku      = []byte("heroku")
)

// "Parse tree" from hell
func serveDrain(w http.ResponseWriter, r *http.Request) {

	var authed bool

	ctx := slog.Context{}
	defer func() { LogWithContext(ctx) }()
	w.Header().Set("Content-Length", "0")

	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		ctx.Count("errors.drain.wrong.method", 1)
		return
	}

	id := r.Header.Get("Logplex-Drain-Token")

	pwd, _ := r.URL.User.Password()
	if r.URL.User.Username() == User && pwd == Password {
		authed = true
	}

	if id == "" && !authed {
		w.WriteHeader(http.StatusBadRequest)
		ctx.Count("errors.drain.missing.token", 1)
		return
	}

	ctx.Count("batch", 1)

	parseStart := time.Now()
	lp := lpx.NewReader(bufio.NewReader(r.Body))

	for lp.Next() {
		ctx.Count("lines.total", 1)
		header := lp.Header()

		// If the syslog App Name Header field containts what looks like a log token,
		// let's assume it's an override of the id and we're getting the data from the magic
		// channel
		if bytes.HasPrefix(header.Name, TokenPrefix) {
			id = string(header.Name)
		}

		// If we still don't have an id, throw an error and try the next line
		if id == "" {
			ctx.Count("errors.token.missing", 1)
			continue
		}

		msg := lp.Bytes()
		switch {
		case bytes.Equal(header.Name, Heroku), bytes.HasPrefix(header.Name, TokenPrefix):
			t, e := time.Parse("2006-01-02T15:04:05.000000+00:00", string(lp.Header().Time))
			if e != nil {
				log.Printf("Error Parsing Time(%s): %q\n", string(lp.Header().Time), e)
				continue
			}
			timestamp := t.UnixNano() / int64(time.Microsecond)

			pid := string(header.Procid)
			switch pid {
			case "router":

				switch {
				// router logs with a H error code in them
				case bytes.Contains(msg, keyCodeH):
					ctx.Count("lines.router.error", 1)
					re := routerError{}
					err := logfmt.Unmarshal(msg, &re)
					if err != nil {
						log.Printf("logfmt unmarshal error: %s\n", err)
						continue
					}
					routerEventPoints <- []interface{}{timestamp, id, re.Code}

				// likely a standard router log
				default:
					ctx.Count("lines.router", 1)
					rm := routerMsg{}
					err := logfmt.Unmarshal(msg, &rm)
					if err != nil {
						log.Printf("logfmt unmarshal error: %s\n", err)
						continue
					}
					routerPoints <- []interface{}{timestamp, id, rm.Status, rm.Service}
				}

				// Non router logs, so either dynos, runtime, etc
			default:
				switch {
				case bytes.HasPrefix(msg, dynoErrorSentinel):
					ctx.Count("lines.dyno.error", 1)
					de, err := parseBytesToDynoError(msg)
					if err != nil {
						log.Printf("Unable to parse dyno error message: %q\n", err)
					}
					dynoEventsPoints <- []interface{}{timestamp, id, string(lp.Header().Procid), "R", de.Code, string(msg)}

				case bytes.Contains(msg, dynoMemMsgSentinel):
					ctx.Count("lines.dyno.mem", 1)
					dm := dynoMemMsg{}
					err := logfmt.Unmarshal(msg, &dm)
					if err != nil {
						log.Printf("logfmt unmarshal error: %s\n", err)
						continue
					}
					if dm.Source != "" {
						dynoMemPoints <- []interface{}{timestamp, id, dm.Source, dm.MemoryCache, dm.MemoryPgpgin, dm.MemoryPgpgout, dm.MemoryRSS, dm.MemorySwap, dm.MemoryTotal}
					}
				case bytes.Contains(msg, dynoLoadMsgSentinel):
					ctx.Count("lines.dyno.load", 1)
					dm := dynoLoadMsg{}
					err := logfmt.Unmarshal(msg, &dm)
					if err != nil {
						log.Printf("logfmt unmarshal error: %s\n", err)
						continue
					}
					if dm.Source != "" {
						dynoLoadPoints <- []interface{}{timestamp, id, dm.Source, dm.LoadAvg1Min, dm.LoadAvg5Min, dm.LoadAvg15Min}
					}
				default: // unknown
					ctx.Count("lines.unknown.heroku", 1)
					if Debug {
						log.Printf("Unknown Heroku Line - Header: PRI: %s, Time: %s, Hostname: %s, Name: %s, ProcId: %s, MsgId: %s - Body: %s",
							header.PrivalVersion,
							header.Time,
							header.Hostname,
							header.Name,
							header.Procid,
							header.Msgid,
							string(msg),
						)
					}
				}
			}
		default: // non heroku lines
			ctx.Count("lines.unknown.user", 1)
			if Debug {
				log.Printf("Unknown User Line - Header: PRI: %s, Time: %s, Hostname: %s, Name: %s, ProcId: %s, MsgId: %s - Body: %s",
					header.PrivalVersion,
					header.Time,
					header.Hostname,
					header.Name,
					header.Procid,
					header.Msgid,
					string(msg),
				)
			}
		}
	}
	ctx.MeasureSince("lines.parse.time", parseStart)

	w.WriteHeader(http.StatusNoContent)
}
<commit_msg>steal from log-iss for now<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"errors"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/bmizerany/lpx"
	"github.com/heroku/slog"
	"github.com/kr/logfmt"
)

var (
	TokenPrefix = []byte("t.")
	Heroku      = []byte("heroku")
)

func checkAuth(r *http.Request) error {
	header := r.Header.Get("Authorization")
	if header == "" {
		return errors.New("Authorization required")
	}
	headerParts := strings.SplitN(header, " ", 2)
	if len(headerParts) != 2 {
		return errors.New("Authorization header is malformed")
	}

	method := headerParts[0]
	if method != "Basic" {
		return errors.New("Only Basic Authorization is accepted")
	}

	encodedUserPass := headerParts[1]
	decodedUserPass, err := base64.StdEncoding.DecodeString(encodedUserPass)
	if err != nil {
		return errors.New("Authorization header is malformed")
	}

	userPassParts := bytes.SplitN(decodedUserPass, []byte{':'}, 2)
	if len(userPassParts) != 2 {
		return errors.New("Authorization header is malformed")
	}

	user := userPassParts[0]
	pass := userPassParts[1]

	if string(user) != User {
		return errors.New("Unknown user")
	}
	if string(pass) != Password {
		return errors.New("Incorrect token")
	}

	return nil
}

// "Parse tree" from hell
func serveDrain(w http.ResponseWriter, r *http.Request) {

	ctx := slog.Context{}
	defer func() { LogWithContext(ctx) }()
	w.Header().Set("Content-Length", "0")

	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		ctx.Count("errors.drain.wrong.method", 1)
		return
	}

	id := r.Header.Get("Logplex-Drain-Token")

	if id == "" {
		if err := checkAuth(r); err != nil {
			w.WriteHeader(http.StatusForbidden)
			ctx.Count("errors.auth.failure", 1)
			return
		}
	}

	ctx.Count("batch", 1)

	parseStart := time.Now()
	lp := lpx.NewReader(bufio.NewReader(r.Body))

	for lp.Next() {
		ctx.Count("lines.total", 1)
		header := lp.Header()

		// If the syslog App Name Header field containts what looks like a log token,
		// let's assume it's an override of the id and we're getting the data from the magic
		// channel
		if bytes.HasPrefix(header.Name, TokenPrefix) {
			id = string(header.Name)
		}

		// If we still don't have an id, throw an error and try the next line
		if id == "" {
			ctx.Count("errors.token.missing", 1)
			continue
		}

		msg := lp.Bytes()
		switch {
		case bytes.Equal(header.Name, Heroku), bytes.HasPrefix(header.Name, TokenPrefix):
			t, e := time.Parse("2006-01-02T15:04:05.000000+00:00", string(lp.Header().Time))
			if e != nil {
				log.Printf("Error Parsing Time(%s): %q\n", string(lp.Header().Time), e)
				continue
			}
			timestamp := t.UnixNano() / int64(time.Microsecond)

			pid := string(header.Procid)
			switch pid {
			case "router":

				switch {
				// router logs with a H error code in them
				case bytes.Contains(msg, keyCodeH):
					ctx.Count("lines.router.error", 1)
					re := routerError{}
					err := logfmt.Unmarshal(msg, &re)
					if err != nil {
						log.Printf("logfmt unmarshal error: %s\n", err)
						continue
					}
					routerEventPoints <- []interface{}{timestamp, id, re.Code}

				// likely a standard router log
				default:
					ctx.Count("lines.router", 1)
					rm := routerMsg{}
					err := logfmt.Unmarshal(msg, &rm)
					if err != nil {
						log.Printf("logfmt unmarshal error: %s\n", err)
						continue
					}
					routerPoints <- []interface{}{timestamp, id, rm.Status, rm.Service}
				}

				// Non router logs, so either dynos, runtime, etc
			default:
				switch {
				case bytes.HasPrefix(msg, dynoErrorSentinel):
					ctx.Count("lines.dyno.error", 1)
					de, err := parseBytesToDynoError(msg)
					if err != nil {
						log.Printf("Unable to parse dyno error message: %q\n", err)
					}
					dynoEventsPoints <- []interface{}{timestamp, id, string(lp.Header().Procid), "R", de.Code, string(msg)}

				case bytes.Contains(msg, dynoMemMsgSentinel):
					ctx.Count("lines.dyno.mem", 1)
					dm := dynoMemMsg{}
					err := logfmt.Unmarshal(msg, &dm)
					if err != nil {
						log.Printf("logfmt unmarshal error: %s\n", err)
						continue
					}
					if dm.Source != "" {
						dynoMemPoints <- []interface{}{timestamp, id, dm.Source, dm.MemoryCache, dm.MemoryPgpgin, dm.MemoryPgpgout, dm.MemoryRSS, dm.MemorySwap, dm.MemoryTotal}
					}
				case bytes.Contains(msg, dynoLoadMsgSentinel):
					ctx.Count("lines.dyno.load", 1)
					dm := dynoLoadMsg{}
					err := logfmt.Unmarshal(msg, &dm)
					if err != nil {
						log.Printf("logfmt unmarshal error: %s\n", err)
						continue
					}
					if dm.Source != "" {
						dynoLoadPoints <- []interface{}{timestamp, id, dm.Source, dm.LoadAvg1Min, dm.LoadAvg5Min, dm.LoadAvg15Min}
					}
				default: // unknown
					ctx.Count("lines.unknown.heroku", 1)
					if Debug {
						log.Printf("Unknown Heroku Line - Header: PRI: %s, Time: %s, Hostname: %s, Name: %s, ProcId: %s, MsgId: %s - Body: %s",
							header.PrivalVersion,
							header.Time,
							header.Hostname,
							header.Name,
							header.Procid,
							header.Msgid,
							string(msg),
						)
					}
				}
			}
		default: // non heroku lines
			ctx.Count("lines.unknown.user", 1)
			if Debug {
				log.Printf("Unknown User Line - Header: PRI: %s, Time: %s, Hostname: %s, Name: %s, ProcId: %s, MsgId: %s - Body: %s",
					header.PrivalVersion,
					header.Time,
					header.Hostname,
					header.Name,
					header.Procid,
					header.Msgid,
					string(msg),
				)
			}
		}
	}
	ctx.MeasureSince("lines.parse.time", parseStart)

	w.WriteHeader(http.StatusNoContent)
}
<|endoftext|>


<commit_before>/**
 * @file session.go
 * @author Mikhail Klementyev jollheef<AT>riseup.net
 * @license GNU AGPLv3
 * @date November, 2015
 * @brief queries for session table
 */

package db

import (
	"database/sql"
	"time"
)

// Session row
type Session struct {
	ID        int
	TeamID    int
	Session   string
	Timestamp time.Time
}

func createSessionTable(db *sql.DB) (err error) {

	_, err = db.Exec(`
	CREATE TABLE IF NOT EXISTS "session" (
		id		SERIAL PRIMARY KEY,
		team_id		INTEGER NOT NULL,
		session		TEXT NOT NULL,
		timestamp	TIMESTAMP with time zone DEFAULT now()
	)`)

	return
}

// AddSession add session and fill id
func AddSession(db *sql.DB, s *Session) (err error) {

	stmt, err := db.Prepare("INSERT INTO session (team_id, session) " +
		"VALUES ($1, $2) RETURNING id")
	if err != nil {
		return
	}

	defer stmt.Close()

	err = stmt.QueryRow(s.TeamID, s.Session).Scan(&s.ID)
	if err != nil {
		return
	}

	return
}

// GetSessionTeam get team id for session
func GetSessionTeam(db *sql.DB, session string) (teamID int, err error) {

	stmt, err := db.Prepare("SELECT team_id FROM session WHERE session=$1")
	if err != nil {
		return
	}

	defer stmt.Close()

	err = stmt.QueryRow(session).Scan(&teamID)
	if err != nil {
		return
	}

	return
}
<commit_msg>Implements get session count<commit_after>/**
 * @file session.go
 * @author Mikhail Klementyev jollheef<AT>riseup.net
 * @license GNU AGPLv3
 * @date November, 2015
 * @brief queries for session table
 */

package db

import (
	"database/sql"
	"time"
)

// Session row
type Session struct {
	ID        int
	TeamID    int
	Session   string
	Timestamp time.Time
}

func createSessionTable(db *sql.DB) (err error) {

	_, err = db.Exec(`
	CREATE TABLE IF NOT EXISTS "session" (
		id		SERIAL PRIMARY KEY,
		team_id		INTEGER NOT NULL,
		session		TEXT NOT NULL,
		timestamp	TIMESTAMP with time zone DEFAULT now()
	)`)

	return
}

// AddSession add session and fill id
func AddSession(db *sql.DB, s *Session) (err error) {

	stmt, err := db.Prepare("INSERT INTO session (team_id, session) " +
		"VALUES ($1, $2) RETURNING id")
	if err != nil {
		return
	}

	defer stmt.Close()

	err = stmt.QueryRow(s.TeamID, s.Session).Scan(&s.ID)
	if err != nil {
		return
	}

	return
}

// GetSessionCount returns count of logged teams for all time
func GetSessionCount(db *sql.DB) (count int, err error) {
	err = db.QueryRow("SELECT COUNT(DISTINCT team_id) FROM session;").Scan(&count)
	return
}

// GetSessionTeam get team id for session
func GetSessionTeam(db *sql.DB, session string) (teamID int, err error) {

	stmt, err := db.Prepare("SELECT team_id FROM session WHERE session=$1")
	if err != nil {
		return
	}

	defer stmt.Close()

	err = stmt.QueryRow(session).Scan(&teamID)
	if err != nil {
		return
	}

	return
}
<|endoftext|>


<commit_before>package dnsd

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/HouzuoGuo/laitos/lalog"
	"github.com/HouzuoGuo/laitos/misc"
	"github.com/HouzuoGuo/laitos/tcpoverdns"
	"github.com/HouzuoGuo/laitos/toolbox"
)

// ProxyRequest is the data sent by a proxy client to initiate a connection
// toward a proxy destination.
type ProxyRequest struct {
	// Network name of the address (e.g. "tcp"), this takes precedence over the
	// port number when the network name is specified.
	Network string `json:"n"`
	// Port number in the absence of network name.
	Port int `json:"p"`
	// Address is the host IP address or network address (IP:port).
	Address string `json:"a"`
	// AccessTOTP is a time-based OTP that authorises the connection request.
	AccessTOTP string `json:"t"`
}

// ProxyConnection consists of a transmission control paired to a TCP connection
// relayed by the transmission control.
type ProxyConnection struct {
	proxy                *Proxy
	tcpConn              *net.TCPConn
	context              context.Context
	tc                   *tcpoverdns.TransmissionControl
	inputSegments        net.Conn
	outputSegmentBacklog []tcpoverdns.Segment
	mutex                *sync.Mutex
	logger               lalog.Logger
}

// Start piping data back and forth between proxy TCP connection and
// transmission control.
// The function blocks until the underlying TC is closed.
func (conn *ProxyConnection) Start() {
	if conn.proxy.Debug {
		conn.logger.Info("Start", "", nil, "starting now")
	}
	beginTimeNano := time.Now().UnixNano()
	defer func() {
		if conn.proxy.Debug {
			conn.logger.Info("Start", "", nil, "closing and lingering")
			conn.tc.DumpState()
		}
		_ = conn.Close()
		misc.TCPOverDNSStats.Trigger(float64(time.Now().UnixNano() - beginTimeNano))
		go func() {
			// Linger a short while before deleting (cease tracking) the
			// connection, as the output segment buffer may still contain
			// useful, final few segments - and crucially it contains the
			// segment with ResetTerminate flag which is expected by the peer
			// transmission control.
			time.Sleep(conn.proxy.Linger)
			conn.proxy.mutex.Lock()
			delete(conn.proxy.connections, conn.tc.ID)
			conn.proxy.mutex.Unlock()
			if conn.proxy.Debug {
				conn.logger.Info("Start", "", nil, "closed and removed from proxy")
			}
		}()
	}()
	// Absorb outgoing segments into the outgoing backlog.
	conn.tc.OutputSegmentCallback = func(seg tcpoverdns.Segment) {
		// Replace the latest keep-alive or ack-only segment (if any), and
		// de-duplicate adjacent identical segments. These measures not only
		// speed up the exchanges but also ensure that peers can communicate
		// properly even if their timing characteristics differ.
		conn.mutex.Lock()
		defer conn.mutex.Unlock()
		var latest tcpoverdns.Segment
		if len(conn.outputSegmentBacklog) > 0 {
			latest = conn.outputSegmentBacklog[len(conn.outputSegmentBacklog)-1]
		}
		if latest.Flags.Has(tcpoverdns.FlagAckOnly) || latest.Flags.Has(tcpoverdns.FlagKeepAlive) {
			// Substitute the ack-only or keep-alive segment with the latest.
			// These segments may contain random data that are not useful.
			if conn.proxy.Debug {
				conn.logger.Info("Start", "", nil, "callback is removing duplicated ack/keepalive segment: %+v", conn.outputSegmentBacklog[len(conn.outputSegmentBacklog)-1])
			}
			conn.outputSegmentBacklog[len(conn.outputSegmentBacklog)-1] = seg
		} else if latest.Equals(seg) {
			// De-duplicate adjacent identical segments.
			if conn.proxy.Debug {
				conn.logger.Info("Start", "", nil, "callback is removing duplicated segment: %+v", seg)
			}
		} else {
			conn.logger.Info("Start", "", nil, "callback is handling segment %+v", seg)
			conn.outputSegmentBacklog = append(conn.outputSegmentBacklog, seg)
		}
	}
	// Carry on with the handshake.
	conn.tc.Start(conn.context)
	conn.tc.WaitState(conn.context, tcpoverdns.StateEstablished)
	if conn.proxy.Debug {
		conn.logger.Info("ProxyConnection.Start", "", nil, "TC is established")
	}
	// Pipe data in both directions.
	if conn.tcpConn != nil {
		// Keep the buffer to minimum to improve responsiveness.
		// The buffer size has nothing to do with segment size.
		go func() {
			if err := misc.PipeConn(conn.logger, false, conn.tc.MaxLifetime, 1, conn.tc, conn.tcpConn); err != nil {
				if conn.proxy.Debug {
					conn.logger.Info("ProxyConnection.Start", "", err, "finished piping from TC to TCP connection")
				}
			}
		}()
		if err := misc.PipeConn(conn.logger, false, conn.tc.MaxLifetime, 1, conn.tcpConn, conn.tc); err != nil {
			if conn.proxy.Debug {
				conn.logger.Info("ProxyConnection.Start", "", err, "finished piping from TCP connection to TC")
			}
		}
	}
	// Wait for the transmission control to close.
	conn.tc.CloseAfterDrained()
	conn.tc.WaitState(conn.context, tcpoverdns.StateClosed)
	// The proxy connection lingers for a short while, see defer.
}

// WaitSegment busy-waits until a new segment is available from the output
// segment backlog, and then pops the segment.
func (conn *ProxyConnection) WaitSegment(ctx context.Context) (tcpoverdns.Segment, bool) {
	for {
		conn.mutex.Lock()
		if len(conn.outputSegmentBacklog) > 0 {
			ret := conn.outputSegmentBacklog[0]
			conn.outputSegmentBacklog = conn.outputSegmentBacklog[1:]
			conn.mutex.Unlock()
			return ret, true
		} else {
			conn.mutex.Unlock()
			select {
			case <-ctx.Done():
				return tcpoverdns.Segment{}, false
			case <-time.After(tcpoverdns.BusyWaitInterval):
				continue
			}
		}
	}
}

// Close and terminate the proxy TCP connection and its transmission control.
func (conn *ProxyConnection) Close() error {
	if conn.tcpConn != nil {
		_ = conn.tcpConn.Close()
	}
	_ = conn.inputSegments.Close()
	_ = conn.tc.Close()
	return nil
}

// Proxy manages the full life cycle of multiple transmission controls created
// for the purpose of relaying TCP connections.
type Proxy struct {
	// DNSDaemon helps the proxy to identify advertising/malware proxy
	// destinations and to refuse serving their clients.
	DNSDaemon *Daemon `json:"-"`

	// MaxReplyLatency is the maximum duration to wait for a reply (outgoing)
	// segment before returning from the Receive function.
	MaxReplyDelay time.Duration `json:"-"`

	// RequestOTPSecret is a TOTP secret for authorising incoming connection
	// requests.
	RequestOTPSecret string `json:"RequestOTPSecret"`
	// Debug enables verbose logging for IO activities.
	Debug bool `json:"Debug"`

	// Linger is a brief period of time for a proxy connection to stay before it
	// is removed from the internal collection of proxy connections. The delay
	// is crucial to allow the final segments of each proxy connection to be
	// received by proxy client - including the segment with ResetTerminate.
	Linger time.Duration `json:"-"`

	// DialTimeout is the timeout used for creating new a proxy TCP connection.
	DialTimeout time.Duration `json:"-"`

	// logger is used to log IO activities when verbose logging is enabled.
	logger lalog.Logger `json:"-"`

	connections map[uint16]*ProxyConnection
	context     context.Context
	cancelFun   func()
	mutex       *sync.Mutex
}

// Start initialises the internal state of the proxy.
func (proxy *Proxy) Start(ctx context.Context) {
	if proxy.MaxReplyDelay == 0 {
		// This default should be greater/longer than transmission control's
		// default AckDelay, or the performance will suffer quite a bit.
		proxy.MaxReplyDelay = 2 * time.Second
	}
	if proxy.DialTimeout == 0 {
		proxy.DialTimeout = 10 * time.Second
	}
	if proxy.Linger == 0 {
		proxy.Linger = 60 * time.Second
	}
	proxy.connections = make(map[uint16]*ProxyConnection)
	proxy.context, proxy.cancelFun = context.WithCancel(ctx)
	proxy.mutex = new(sync.Mutex)
	proxy.logger = lalog.Logger{ComponentName: "TCProxy"}
}

// Receive processes an incoming segment and relay the segment to an existing
// transmission control, or create a new transmission control for the proxy
// destination.
func (proxy *Proxy) Receive(in tcpoverdns.Segment) (tcpoverdns.Segment, bool) {
	proxy.mutex.Lock()
	conn, exists := proxy.connections[in.ID]
	proxy.mutex.Unlock()
	if !exists {
		// Connect to the proxy destination.
		var req ProxyRequest
		proxy.logger.Info("Receive", in.ID, nil, "new connection request - seg: %+v, req: %+v", in, req)
		if len(in.Data) < tcpoverdns.InitiatorConfigLen {
			proxy.logger.Warning("Receive", in.ID, nil, "received a malformed segment possibly from a stale TC")
			return tcpoverdns.Segment{}, false
		}
		if err := json.Unmarshal(in.Data[tcpoverdns.InitiatorConfigLen:], &req); err != nil {
			proxy.logger.Warning("Receive", in.ID, err, "failed to deserialise proxy request")
			return tcpoverdns.Segment{}, false
		}
		if proxy.DNSDaemon != nil && proxy.DNSDaemon.IsInBlacklist(req.Address) {
			proxy.logger.Info("Receive", in.ID, nil, "refusing connection to blacklisted destination %q", req.Address)
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		prev, curr, next, err := toolbox.GetTwoFACodes(proxy.RequestOTPSecret)
		if err != nil {
			proxy.logger.Info("Receive", in.ID, nil, "failed to calculate TOTP codes")
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		if req.AccessTOTP != prev && req.AccessTOTP != curr && req.AccessTOTP != next {
			proxy.logger.Warning("Receive", in.ID, nil, "the request failed OTP check")
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		// Construct the transmission control at proxy's side.
		proxyIn, tcIn := net.Pipe()
		tc := &tcpoverdns.TransmissionControl{
			Debug:  proxy.Debug,
			LogTag: "ProxyServer",
			ID:     in.ID,
			// This transmission control is a responder during the handshake.
			Initiator:      false,
			InputTransport: tcIn,
			MaxLifetime:    30 * time.Minute,
			// In practice there are occasionally bursts of tens of errors at a
			// time before recovery.
			MaxTransportErrors: 200,
			// The duration of all retransmissions (if all go unacknowledged) is
			// MaxRetransmissions x SlidingWindowWaitDuration.
			MaxRetransmissions: 200,
			// The output transport is not used. Instead, the output segments
			// are kept in a backlog.
			OutputTransport: io.Discard,
			// The segment length and sliding window length are set by the
			// initiator using InitiatorConfig.
		}
		// Connect to the intended destination.
		var dialNet, dialDest string
		if req.Network == "" {
			dialNet = "tcp"
			dialDest = fmt.Sprintf("%s:%d", req.Address, req.Port)
		} else {
			dialNet = req.Network
			dialDest = req.Address
		}
		netConn, err := net.DialTimeout(dialNet, dialDest, proxy.DialTimeout)
		if err != nil {
			// Immediately close the transmission control if the destination is
			// unreachable.
			proxy.logger.Warning("Receive", in.ID, err, "failed to connect to proxy destination %s %s", dialNet, dialDest)
			// Proceed with handshake, but there will be no data coming through
			// the transmission control and it will be closed shortly.
		}
		var tcpConn *net.TCPConn
		if netConn != nil {
			tcpConn = netConn.(*net.TCPConn)
			misc.TweakTCPConnection(tcpConn, tc.MaxLifetime)
		}
		// Track the new proxy connection.
		conn = &ProxyConnection{
			proxy:         proxy,
			tcpConn:       tcpConn,
			context:       proxy.context,
			tc:            tc,
			inputSegments: proxyIn,
			mutex:         new(sync.Mutex),
			logger: lalog.Logger{
				ComponentName: "ProxyServer",
				ComponentID: []lalog.LoggerIDField{
					{Key: "TCID", Value: tc.ID},
				},
			},
		}
		proxy.mutex.Lock()
		proxy.connections[in.ID] = conn
		proxy.mutex.Unlock()
		// Go ahead with handshake and data transfer.
		go conn.Start()
	}
	if _, err := conn.inputSegments.Write(in.Packet()); err != nil {
		_ = conn.Close()
	}
	if proxy.Debug {
		proxy.logger.Info("Receive", in.ID, nil, "waiting for a reply outbound segment")
	}
	waitCtx, cancel := context.WithTimeout(proxy.context, proxy.MaxReplyDelay)
	defer cancel()
	seg, hasSeg := conn.WaitSegment(waitCtx)
	return seg, hasSeg
}

// Close terminates all ongoing transmission controls.
// The function always returns nil.
func (proxy *Proxy) Close() error {
	proxy.mutex.Lock()
	// Terminate all TCs.
	for _, conn := range proxy.connections {
		_ = conn.Close()
	}
	defer proxy.mutex.Unlock()
	// Terminate all proxy connections.
	proxy.cancelFun()
	return nil
}
<commit_msg>dnsd: hide tcp-over-dns callback debug message behind the debug switch<commit_after>package dnsd

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/HouzuoGuo/laitos/lalog"
	"github.com/HouzuoGuo/laitos/misc"
	"github.com/HouzuoGuo/laitos/tcpoverdns"
	"github.com/HouzuoGuo/laitos/toolbox"
)

// ProxyRequest is the data sent by a proxy client to initiate a connection
// toward a proxy destination.
type ProxyRequest struct {
	// Network name of the address (e.g. "tcp"), this takes precedence over the
	// port number when the network name is specified.
	Network string `json:"n"`
	// Port number in the absence of network name.
	Port int `json:"p"`
	// Address is the host IP address or network address (IP:port).
	Address string `json:"a"`
	// AccessTOTP is a time-based OTP that authorises the connection request.
	AccessTOTP string `json:"t"`
}

// ProxyConnection consists of a transmission control paired to a TCP connection
// relayed by the transmission control.
type ProxyConnection struct {
	proxy                *Proxy
	tcpConn              *net.TCPConn
	context              context.Context
	tc                   *tcpoverdns.TransmissionControl
	inputSegments        net.Conn
	outputSegmentBacklog []tcpoverdns.Segment
	mutex                *sync.Mutex
	logger               lalog.Logger
}

// Start piping data back and forth between proxy TCP connection and
// transmission control.
// The function blocks until the underlying TC is closed.
func (conn *ProxyConnection) Start() {
	if conn.proxy.Debug {
		conn.logger.Info("Start", "", nil, "starting now")
	}
	beginTimeNano := time.Now().UnixNano()
	defer func() {
		if conn.proxy.Debug {
			conn.logger.Info("Start", "", nil, "closing and lingering")
			conn.tc.DumpState()
		}
		_ = conn.Close()
		misc.TCPOverDNSStats.Trigger(float64(time.Now().UnixNano() - beginTimeNano))
		go func() {
			// Linger a short while before deleting (cease tracking) the
			// connection, as the output segment buffer may still contain
			// useful, final few segments - and crucially it contains the
			// segment with ResetTerminate flag which is expected by the peer
			// transmission control.
			time.Sleep(conn.proxy.Linger)
			conn.proxy.mutex.Lock()
			delete(conn.proxy.connections, conn.tc.ID)
			conn.proxy.mutex.Unlock()
			if conn.proxy.Debug {
				conn.logger.Info("Start", "", nil, "closed and removed from proxy")
			}
		}()
	}()
	// Absorb outgoing segments into the outgoing backlog.
	conn.tc.OutputSegmentCallback = func(seg tcpoverdns.Segment) {
		// Replace the latest keep-alive or ack-only segment (if any), and
		// de-duplicate adjacent identical segments. These measures not only
		// speed up the exchanges but also ensure that peers can communicate
		// properly even if their timing characteristics differ.
		conn.mutex.Lock()
		defer conn.mutex.Unlock()
		var latest tcpoverdns.Segment
		if len(conn.outputSegmentBacklog) > 0 {
			latest = conn.outputSegmentBacklog[len(conn.outputSegmentBacklog)-1]
		}
		if latest.Flags.Has(tcpoverdns.FlagAckOnly) || latest.Flags.Has(tcpoverdns.FlagKeepAlive) {
			// Substitute the ack-only or keep-alive segment with the latest.
			// These segments may contain random data that are not useful.
			if conn.proxy.Debug {
				conn.logger.Info("Start", "", nil, "callback is removing duplicated ack/keepalive segment: %+v", conn.outputSegmentBacklog[len(conn.outputSegmentBacklog)-1])
			}
			conn.outputSegmentBacklog[len(conn.outputSegmentBacklog)-1] = seg
		} else if latest.Equals(seg) {
			// De-duplicate adjacent identical segments.
			if conn.proxy.Debug {
				conn.logger.Info("Start", "", nil, "callback is removing duplicated segment: %+v", seg)
			}
		} else {
			if conn.proxy.Debug {
				conn.logger.Info("Start", "", nil, "callback is handling segment %+v", seg)
			}
			conn.outputSegmentBacklog = append(conn.outputSegmentBacklog, seg)
		}
	}
	// Carry on with the handshake.
	conn.tc.Start(conn.context)
	conn.tc.WaitState(conn.context, tcpoverdns.StateEstablished)
	if conn.proxy.Debug {
		conn.logger.Info("ProxyConnection.Start", "", nil, "TC is established")
	}
	// Pipe data in both directions.
	if conn.tcpConn != nil {
		// Keep the buffer to minimum to improve responsiveness.
		// The buffer size has nothing to do with segment size.
		go func() {
			if err := misc.PipeConn(conn.logger, false, conn.tc.MaxLifetime, 1, conn.tc, conn.tcpConn); err != nil {
				if conn.proxy.Debug {
					conn.logger.Info("ProxyConnection.Start", "", err, "finished piping from TC to TCP connection")
				}
			}
		}()
		if err := misc.PipeConn(conn.logger, false, conn.tc.MaxLifetime, 1, conn.tcpConn, conn.tc); err != nil {
			if conn.proxy.Debug {
				conn.logger.Info("ProxyConnection.Start", "", err, "finished piping from TCP connection to TC")
			}
		}
	}
	// Wait for the transmission control to close.
	conn.tc.CloseAfterDrained()
	conn.tc.WaitState(conn.context, tcpoverdns.StateClosed)
	// The proxy connection lingers for a short while, see defer.
}

// WaitSegment busy-waits until a new segment is available from the output
// segment backlog, and then pops the segment.
func (conn *ProxyConnection) WaitSegment(ctx context.Context) (tcpoverdns.Segment, bool) {
	for {
		conn.mutex.Lock()
		if len(conn.outputSegmentBacklog) > 0 {
			ret := conn.outputSegmentBacklog[0]
			conn.outputSegmentBacklog = conn.outputSegmentBacklog[1:]
			conn.mutex.Unlock()
			return ret, true
		} else {
			conn.mutex.Unlock()
			select {
			case <-ctx.Done():
				return tcpoverdns.Segment{}, false
			case <-time.After(tcpoverdns.BusyWaitInterval):
				continue
			}
		}
	}
}

// Close and terminate the proxy TCP connection and its transmission control.
func (conn *ProxyConnection) Close() error {
	if conn.tcpConn != nil {
		_ = conn.tcpConn.Close()
	}
	_ = conn.inputSegments.Close()
	_ = conn.tc.Close()
	return nil
}

// Proxy manages the full life cycle of multiple transmission controls created
// for the purpose of relaying TCP connections.
type Proxy struct {
	// DNSDaemon helps the proxy to identify advertising/malware proxy
	// destinations and to refuse serving their clients.
	DNSDaemon *Daemon `json:"-"`

	// MaxReplyLatency is the maximum duration to wait for a reply (outgoing)
	// segment before returning from the Receive function.
	MaxReplyDelay time.Duration `json:"-"`

	// RequestOTPSecret is a TOTP secret for authorising incoming connection
	// requests.
	RequestOTPSecret string `json:"RequestOTPSecret"`
	// Debug enables verbose logging for IO activities.
	Debug bool `json:"Debug"`

	// Linger is a brief period of time for a proxy connection to stay before it
	// is removed from the internal collection of proxy connections. The delay
	// is crucial to allow the final segments of each proxy connection to be
	// received by proxy client - including the segment with ResetTerminate.
	Linger time.Duration `json:"-"`

	// DialTimeout is the timeout used for creating new a proxy TCP connection.
	DialTimeout time.Duration `json:"-"`

	// logger is used to log IO activities when verbose logging is enabled.
	logger lalog.Logger `json:"-"`

	connections map[uint16]*ProxyConnection
	context     context.Context
	cancelFun   func()
	mutex       *sync.Mutex
}

// Start initialises the internal state of the proxy.
func (proxy *Proxy) Start(ctx context.Context) {
	if proxy.MaxReplyDelay == 0 {
		// This default should be greater/longer than transmission control's
		// default AckDelay, or the performance will suffer quite a bit.
		proxy.MaxReplyDelay = 2 * time.Second
	}
	if proxy.DialTimeout == 0 {
		proxy.DialTimeout = 10 * time.Second
	}
	if proxy.Linger == 0 {
		proxy.Linger = 60 * time.Second
	}
	proxy.connections = make(map[uint16]*ProxyConnection)
	proxy.context, proxy.cancelFun = context.WithCancel(ctx)
	proxy.mutex = new(sync.Mutex)
	proxy.logger = lalog.Logger{ComponentName: "TCProxy"}
}

// Receive processes an incoming segment and relay the segment to an existing
// transmission control, or create a new transmission control for the proxy
// destination.
func (proxy *Proxy) Receive(in tcpoverdns.Segment) (tcpoverdns.Segment, bool) {
	proxy.mutex.Lock()
	conn, exists := proxy.connections[in.ID]
	proxy.mutex.Unlock()
	if !exists {
		// Connect to the proxy destination.
		var req ProxyRequest
		proxy.logger.Info("Receive", in.ID, nil, "new connection request - seg: %+v, req: %+v", in, req)
		if len(in.Data) < tcpoverdns.InitiatorConfigLen {
			proxy.logger.Warning("Receive", in.ID, nil, "received a malformed segment possibly from a stale TC")
			return tcpoverdns.Segment{}, false
		}
		if err := json.Unmarshal(in.Data[tcpoverdns.InitiatorConfigLen:], &req); err != nil {
			proxy.logger.Warning("Receive", in.ID, err, "failed to deserialise proxy request")
			return tcpoverdns.Segment{}, false
		}
		if proxy.DNSDaemon != nil && proxy.DNSDaemon.IsInBlacklist(req.Address) {
			proxy.logger.Info("Receive", in.ID, nil, "refusing connection to blacklisted destination %q", req.Address)
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		prev, curr, next, err := toolbox.GetTwoFACodes(proxy.RequestOTPSecret)
		if err != nil {
			proxy.logger.Info("Receive", in.ID, nil, "failed to calculate TOTP codes")
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		if req.AccessTOTP != prev && req.AccessTOTP != curr && req.AccessTOTP != next {
			proxy.logger.Warning("Receive", in.ID, nil, "the request failed OTP check")
			return tcpoverdns.Segment{ID: in.ID, Flags: tcpoverdns.FlagReset}, true
		}
		// Construct the transmission control at proxy's side.
		proxyIn, tcIn := net.Pipe()
		tc := &tcpoverdns.TransmissionControl{
			Debug:  proxy.Debug,
			LogTag: "ProxyServer",
			ID:     in.ID,
			// This transmission control is a responder during the handshake.
			Initiator:      false,
			InputTransport: tcIn,
			MaxLifetime:    30 * time.Minute,
			// In practice there are occasionally bursts of tens of errors at a
			// time before recovery.
			MaxTransportErrors: 200,
			// The duration of all retransmissions (if all go unacknowledged) is
			// MaxRetransmissions x SlidingWindowWaitDuration.
			MaxRetransmissions: 200,
			// The output transport is not used. Instead, the output segments
			// are kept in a backlog.
			OutputTransport: io.Discard,
			// The segment length and sliding window length are set by the
			// initiator using InitiatorConfig.
		}
		// Connect to the intended destination.
		var dialNet, dialDest string
		if req.Network == "" {
			dialNet = "tcp"
			dialDest = fmt.Sprintf("%s:%d", req.Address, req.Port)
		} else {
			dialNet = req.Network
			dialDest = req.Address
		}
		netConn, err := net.DialTimeout(dialNet, dialDest, proxy.DialTimeout)
		if err != nil {
			// Immediately close the transmission control if the destination is
			// unreachable.
			proxy.logger.Warning("Receive", in.ID, err, "failed to connect to proxy destination %s %s", dialNet, dialDest)
			// Proceed with handshake, but there will be no data coming through
			// the transmission control and it will be closed shortly.
		}
		var tcpConn *net.TCPConn
		if netConn != nil {
			tcpConn = netConn.(*net.TCPConn)
			misc.TweakTCPConnection(tcpConn, tc.MaxLifetime)
		}
		// Track the new proxy connection.
		conn = &ProxyConnection{
			proxy:         proxy,
			tcpConn:       tcpConn,
			context:       proxy.context,
			tc:            tc,
			inputSegments: proxyIn,
			mutex:         new(sync.Mutex),
			logger: lalog.Logger{
				ComponentName: "ProxyServer",
				ComponentID: []lalog.LoggerIDField{
					{Key: "TCID", Value: tc.ID},
				},
			},
		}
		proxy.mutex.Lock()
		proxy.connections[in.ID] = conn
		proxy.mutex.Unlock()
		// Go ahead with handshake and data transfer.
		go conn.Start()
	}
	if _, err := conn.inputSegments.Write(in.Packet()); err != nil {
		_ = conn.Close()
	}
	if proxy.Debug {
		proxy.logger.Info("Receive", in.ID, nil, "waiting for a reply outbound segment")
	}
	waitCtx, cancel := context.WithTimeout(proxy.context, proxy.MaxReplyDelay)
	defer cancel()
	seg, hasSeg := conn.WaitSegment(waitCtx)
	return seg, hasSeg
}

// Close terminates all ongoing transmission controls.
// The function always returns nil.
func (proxy *Proxy) Close() error {
	proxy.mutex.Lock()
	// Terminate all TCs.
	for _, conn := range proxy.connections {
		_ = conn.Close()
	}
	defer proxy.mutex.Unlock()
	// Terminate all proxy connections.
	proxy.cancelFun()
	return nil
}
<|endoftext|>


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

// File I/O for logs.

package glog

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strings"
	"time"
)

// MaxSize is the maximum size of a log file in bytes.
var MaxSize uint64 = 1024 * 1024 * 1800

// If non-empty, overrides the choice of directory in which to write logs.
// See createLogDirs for the full list of possible destinations.
var logDir = flag.String("log_dir", "", "If non-empty, write log files in this directory")

var (
	pid      = os.Getpid()
	program  = filepath.Base(os.Args[0])
	host     = "unknownhost"
	userName = "unknownuser"
)

func init() {
	h, err := os.Hostname()
	if err == nil {
		host = shortHostname(h)
	}

	current, err := user.Current()
	if err == nil {
		userName = current.Username
	}

	// Sanitize userName since it may contain filepath separators on Windows.
	userName = strings.Replace(userName, `\`, "_", -1)
}

// shortHostname returns its argument, truncating at the first period.
// For instance, given "www.google.com" it returns "www".
func shortHostname(hostname string) string {
	if i := strings.Index(hostname, "."); i >= 0 {
		return hostname[:i]
	}
	return hostname
}

// logName returns a new log file name containing tag, with start time t, and
// the name for the symlink for tag.
func logName(tag string, t time.Time) (name, link string) {
	name = fmt.Sprintf("%s.%s.%s.log.%s.%04d%02d%02d-%02d%02d%02d.%d",
		program,
		host,
		userName,
		tag,
		t.Year(),
		t.Month(),
		t.Day(),
		t.Hour(),
		t.Minute(),
		t.Second(),
		pid)
	return name, program + "." + tag
}

// create creates a new log file and returns the file and its filename, which
// contains tag ("INFO", "FATAL", etc.) and t.  If the file is created
// successfully, create also attempts to update the symlink for that tag, ignoring
// errors.
func create(tag string, t time.Time, lrc bool) (f *os.File, filename string, err error) {
	if *logDir == "" {
		return nil, "", errors.New("log: no log dir")
	}
	name, link := logName(tag, t)
	var lastErr error
	if !lrc {
		fname := filepath.Join(*logDir, name)
		f, err := os.Create(fname)
		if err == nil {
			symlink := filepath.Join(*logDir, link)
			os.Remove(symlink)        // ignore err
			os.Symlink(name, symlink) // ignore err
			return f, fname, nil
		}
		lastErr = err
	} else {
		fname := filepath.Join(*logDir, link)
		f, err := os.OpenFile(fname, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0660)
		if err == nil {
			return f, fname, nil
		}
		lastErr = err
	}
	return nil, "", fmt.Errorf("log: cannot create log: %v", lastErr)
}
<commit_msg>use . as default logDir<commit_after>// Go support for leveled logs, analogous to https://code.google.com/p/google-glog/
//
// Copyright 2013 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// File I/O for logs.

package glog

import (
	"flag"
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"strings"
	"time"
)

// MaxSize is the maximum size of a log file in bytes.
var MaxSize uint64 = 1024 * 1024 * 1800

// If non-empty, overrides the choice of directory in which to write logs.
// See createLogDirs for the full list of possible destinations.
var logDir = flag.String("log_dir", ".", "If non-empty, write log files in this directory")

var (
	pid      = os.Getpid()
	program  = filepath.Base(os.Args[0])
	host     = "unknownhost"
	userName = "unknownuser"
)

func init() {
	h, err := os.Hostname()
	if err == nil {
		host = shortHostname(h)
	}

	current, err := user.Current()
	if err == nil {
		userName = current.Username
	}

	// Sanitize userName since it may contain filepath separators on Windows.
	userName = strings.Replace(userName, `\`, "_", -1)
}

// shortHostname returns its argument, truncating at the first period.
// For instance, given "www.google.com" it returns "www".
func shortHostname(hostname string) string {
	if i := strings.Index(hostname, "."); i >= 0 {
		return hostname[:i]
	}
	return hostname
}

// logName returns a new log file name containing tag, with start time t, and
// the name for the symlink for tag.
func logName(tag string, t time.Time) (name, link string) {
	name = fmt.Sprintf("%s.%s.%s.log.%s.%04d%02d%02d-%02d%02d%02d.%d",
		program,
		host,
		userName,
		tag,
		t.Year(),
		t.Month(),
		t.Day(),
		t.Hour(),
		t.Minute(),
		t.Second(),
		pid)
	return name, program + "." + tag
}

// create creates a new log file and returns the file and its filename, which
// contains tag ("INFO", "FATAL", etc.) and t.  If the file is created
// successfully, create also attempts to update the symlink for that tag, ignoring
// errors.
func create(tag string, t time.Time, lrc bool) (f *os.File, filename string, err error) {
	name, link := logName(tag, t)
	var lastErr error
	if !lrc {
		fname := filepath.Join(*logDir, name)
		f, err := os.Create(fname)
		if err == nil {
			symlink := filepath.Join(*logDir, link)
			os.Remove(symlink)        // ignore err
			os.Symlink(name, symlink) // ignore err
			return f, fname, nil
		}
		lastErr = err
	} else {
		fname := filepath.Join(*logDir, link)
		f, err := os.OpenFile(fname, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0660)
		if err == nil {
			return f, fname, nil
		}
		lastErr = err
	}
	return nil, "", fmt.Errorf("log: cannot create log: %v", lastErr)
}
<|endoftext|>


<commit_before>package antlr4

import (
	"fmt"
)

//
// Specialized {@link Set}{@code <}{@link ATNConfig}{@code >} that can track
// info about the set, with support for combining similar configurations using a
// graph-structured stack.
///

func hashATNConfig(c interface{}) string {
	return c.(*ATNConfig).shortHashString()
}

func equalATNConfigs(a, b interface{}) bool {
	if a == b {
		return true
	}
	if a == nil || b == nil {
		return false
	}
	return a.(*ATNConfig).state.GetStateNumber() == b.(*ATNConfig).state.GetStateNumber() &&
		a.(*ATNConfig).alt == b.(*ATNConfig).alt &&
		a.(*ATNConfig).semanticContext.equals(b.(*ATNConfig).semanticContext)
}

type ATNConfigSet struct {
	readOnly             bool
	fullCtx              bool
	configLookup         *Set
	conflictingAlts      *BitSet
	cachedHashString     string
	hasSemanticContext   bool
	dipsIntoOuterContext bool
	configs              []IATNConfig
	uniqueAlt            int
}

func NewATNConfigSet(fullCtx bool) *ATNConfigSet {

	a := new(ATNConfigSet)

	// The reason that we need a.is because we don't want the hash map to use
	// the standard hash code and equals. We need all configurations with the
	// same
	// {@code (s,i,_,semctx)} to be equal. Unfortunately, a.key effectively
	// doubles
	// the number of objects associated with ATNConfigs. The other solution is
	// to
	// use a hash table that lets us specify the equals/hashcode operation.
	// All configs but hashed by (s, i, _, pi) not including context. Wiped out
	// when we go readonly as a.set becomes a DFA state.
	a.configLookup = NewSet(hashATNConfig, equalATNConfigs)
	// Indicates that a.configuration set is part of a full context
	// LL prediction. It will be used to determine how to merge $. With SLL
	// it's a wildcard whereas it is not for LL context merge.
	a.fullCtx = fullCtx
	// Indicates that the set of configurations is read-only. Do not
	// allow any code to manipulate the set DFA states will point at
	// the sets and they must not change. a.does not protect the other
	// fields in particular, conflictingAlts is set after
	// we've made a.readonly.
	a.readOnly = false
	// Track the elements as they are added to the set supports Get(i)///
	a.configs = make([]IATNConfig, 0)

	// TODO: these fields make me pretty uncomfortable but nice to pack up info
	// together, saves recomputation
	// TODO: can we track conflicts as they are added to save scanning configs
	// later?
	a.uniqueAlt = 0
	a.conflictingAlts = nil

	// Used in parser and lexer. In lexer, it indicates we hit a pred
	// while computing a closure operation. Don't make a DFA state from a.
	a.hasSemanticContext = false
	a.dipsIntoOuterContext = false

	a.cachedHashString = "-1"

	return a
}

// Adding a Newconfig means merging contexts with existing configs for
// {@code (s, i, pi, _)}, where {@code s} is the
// {@link ATNConfig//state}, {@code i} is the {@link ATNConfig//alt}, and
// {@code pi} is the {@link ATNConfig//semanticContext}. We use
// {@code (s,i,pi)} as key.
//
//
This method updates {@link //dipsIntoOuterContext} and
// {@link //hasSemanticContext} when necessary.

// /
func (this *ATNConfigSet) add(config IATNConfig, mergeCache *DoubleDict) bool {

	if this.readOnly {
		panic("This set is readonly")
	}
	if config.GetSemanticContext() != SemanticContextNONE {
		this.hasSemanticContext = true
	}
	if config.GetReachesIntoOuterContext() > 0 {
		this.dipsIntoOuterContext = true
	}
	var existing = this.configLookup.add(config).(IATNConfig)
	if existing == config {
		this.cachedHashString = "-1"
		this.configs = append(this.configs, config) // track order here
		return true
	}
	// a previous (s,i,pi,_), merge with it and save result
	var rootIsWildcard = !this.fullCtx
	var merged = merge(existing.GetContext(), config.GetContext(), rootIsWildcard, mergeCache)
	// no need to check for existing.context, config.context in cache
	// since only way to create Newgraphs is "call rule" and here. We
	// cache at both places.
	existing.SetReachesIntoOuterContext(intMax(existing.GetReachesIntoOuterContext(), config.GetReachesIntoOuterContext()))
	// make sure to preserve the precedence filter suppression during the merge
	if config.getPrecedenceFilterSuppressed() {
		existing.setPrecedenceFilterSuppressed(true)
	}
	existing.SetContext(merged) // replace context no need to alt mapping
	return true
}

func (this *ATNConfigSet) GetStates() *Set {
	var states = NewSet(nil, nil)
	for i := 0; i < len(this.configs); i++ {
		states.add(this.configs[i].GetState())
	}
	return states
}

func (this *ATNConfigSet) getPredicates() []SemanticContext {
	var preds = make([]SemanticContext, 0)
	for i := 0; i < len(this.configs); i++ {
		c := this.configs[i].GetSemanticContext()
		if c != SemanticContextNONE {
			preds = append(preds, c)
		}
	}
	return preds
}

func (this *ATNConfigSet) getItems() []IATNConfig {
	return this.configs
}

func (this *ATNConfigSet) optimizeConfigs(interpreter *ATNSimulator) {
	if this.readOnly {
		panic("This set is readonly")
	}
	if this.configLookup.length() == 0 {
		return
	}
	for i := 0; i < len(this.configs); i++ {
		var config = this.configs[i]
		config.SetContext(interpreter.getCachedContext(config.GetContext()))
	}
}

func (this *ATNConfigSet) addAll(coll []*ATNConfig) bool {
	for i := 0; i < len(coll); i++ {
		this.add(coll[i], nil)
	}
	return false
}

func (this *ATNConfigSet) equals(other interface{}) bool {
	if this == other {
		return true
	} else if _, ok := other.(*ATNConfigSet); !ok {
		return false
	}

	other2 := other.(*ATNConfigSet)

	return this.configs != nil &&
		//			this.configs.equals(other2.configs) && // TODO is this necessary?
		this.fullCtx == other2.fullCtx &&
		this.uniqueAlt == other2.uniqueAlt &&
		this.conflictingAlts == other2.conflictingAlts &&
		this.hasSemanticContext == other2.hasSemanticContext &&
		this.dipsIntoOuterContext == other2.dipsIntoOuterContext
}

func (this *ATNConfigSet) hashString() string {
	if this.readOnly {
		if this.cachedHashString == "-1" {
			this.cachedHashString = this.hashConfigs()
		}
		return this.cachedHashString
	} else {
		return this.hashConfigs()
	}
}

func (this *ATNConfigSet) hashConfigs() string {
	var s = ""
	for _, c := range this.configs {
		s += fmt.Sprint(c)
	}
	return s
}

func (this *ATNConfigSet) length() int {
	return len(this.configs)
}

func (this *ATNConfigSet) isEmpty() bool {
	return len(this.configs) == 0
}

func (this *ATNConfigSet) contains(item *ATNConfig) bool {
	if this.configLookup == nil {
		panic("This method is not implemented for readonly sets.")
	}
	return this.configLookup.contains(item)
}

func (this *ATNConfigSet) containsFast(item *ATNConfig) bool {
	if this.configLookup == nil {
		panic("This method is not implemented for readonly sets.")
	}
	return this.configLookup.contains(item) // TODO containsFast is not implemented for Set
}

func (this *ATNConfigSet) clear() {
	if this.readOnly {
		panic("This set is readonly")
	}
	this.configs = make([]IATNConfig, 0)
	this.cachedHashString = "-1"
	this.configLookup = NewSet(hashATNConfig, equalATNConfigs)
}

func (this *ATNConfigSet) setReadonly(readOnly bool) {
	this.readOnly = readOnly
	if readOnly {
		this.configLookup = nil // can't mod, no need for lookup cache
	}
}

func (this *ATNConfigSet) toString() string {
	return ""

	return Utils.arrayToString(this.configs) +
		(this.hasSemanticContext ? ",hasSemanticContext=" + this.hasSemanticContext : "") +
		(this.uniqueAlt != ATN.INVALID_ALT_NUMBER ? ",uniqueAlt=" + this.uniqueAlt : "") +
		(this.conflictingAlts != nil ? ",conflictingAlts=" + this.conflictingAlts : "") +
		(this.dipsIntoOuterContext ? ",dipsIntoOuterContext" : "")
}

type OrderedATNConfigSet struct {
	*ATNConfigSet
}

func NewOrderedATNConfigSet() *OrderedATNConfigSet {

	this := new(OrderedATNConfigSet)

	this.ATNConfigSet = NewATNConfigSet(false)
	this.configLookup = NewSet(nil, nil)

	return this
}
<commit_msg>Impl for toString<commit_after>package antlr4

import (
	"fmt"
)

//
// Specialized {@link Set}{@code <}{@link ATNConfig}{@code >} that can track
// info about the set, with support for combining similar configurations using a
// graph-structured stack.
///

func hashATNConfig(c interface{}) string {
	return c.(*ATNConfig).shortHashString()
}

func equalATNConfigs(a, b interface{}) bool {
	if a == b {
		return true
	}
	if a == nil || b == nil {
		return false
	}

	ai,ok := a.(IATNConfig)
	bi,ok1 := b.(IATNConfig)

	if (!ok || !ok1) {
		return false
	}

	return ai.GetState().GetStateNumber() == bi.GetState().GetStateNumber() &&
		ai.GetAlt() == bi.GetAlt() &&
		ai.GetSemanticContext().equals(bi.GetSemanticContext())
}

type ATNConfigSet struct {
	readOnly             bool
	fullCtx              bool
	configLookup         *Set
	conflictingAlts      *BitSet
	cachedHashString     string
	hasSemanticContext   bool``
	dipsIntoOuterContext bool
	configs              []IATNConfig
	uniqueAlt            int
}

func NewATNConfigSet(fullCtx bool) *ATNConfigSet {

	a := new(ATNConfigSet)

	// The reason that we need a.is because we don't want the hash map to use
	// the standard hash code and equals. We need all configurations with the
	// same
	// {@code (s,i,_,semctx)} to be equal. Unfortunately, a.key effectively
	// doubles
	// the number of objects associated with ATNConfigs. The other solution is
	// to
	// use a hash table that lets us specify the equals/hashcode operation.
	// All configs but hashed by (s, i, _, pi) not including context. Wiped out
	// when we go readonly as a.set becomes a DFA state.
	a.configLookup = NewSet(hashATNConfig, equalATNConfigs)
	// Indicates that a.configuration set is part of a full context
	// LL prediction. It will be used to determine how to merge $. With SLL
	// it's a wildcard whereas it is not for LL context merge.
	a.fullCtx = fullCtx
	// Indicates that the set of configurations is read-only. Do not
	// allow any code to manipulate the set DFA states will point at
	// the sets and they must not change. a.does not protect the other
	// fields in particular, conflictingAlts is set after
	// we've made a.readonly.
	a.readOnly = false
	// Track the elements as they are added to the set supports Get(i)///
	a.configs = make([]IATNConfig, 0)

	// TODO: these fields make me pretty uncomfortable but nice to pack up info
	// together, saves recomputation
	// TODO: can we track conflicts as they are added to save scanning configs
	// later?
	a.uniqueAlt = 0
	a.conflictingAlts = nil

	// Used in parser and lexer. In lexer, it indicates we hit a pred
	// while computing a closure operation. Don't make a DFA state from a.
	a.hasSemanticContext = false
	a.dipsIntoOuterContext = false

	a.cachedHashString = "-1"

	return a
}

// Adding a Newconfig means merging contexts with existing configs for
// {@code (s, i, pi, _)}, where {@code s} is the
// {@link ATNConfig//state}, {@code i} is the {@link ATNConfig//alt}, and
// {@code pi} is the {@link ATNConfig//semanticContext}. We use
// {@code (s,i,pi)} as key.
//
//
This method updates {@link //dipsIntoOuterContext} and
// {@link //hasSemanticContext} when necessary.

// /
func (this *ATNConfigSet) add(config IATNConfig, mergeCache *DoubleDict) bool {

	if this.readOnly {
		panic("This set is readonly")
	}
	if config.GetSemanticContext() != SemanticContextNONE {
		this.hasSemanticContext = true
	}
	if config.GetReachesIntoOuterContext() > 0 {
		this.dipsIntoOuterContext = true
	}
	var existing = this.configLookup.add(config).(IATNConfig)
	if existing == config {
		this.cachedHashString = "-1"
		this.configs = append(this.configs, config) // track order here
		return true
	}
	// a previous (s,i,pi,_), merge with it and save result
	var rootIsWildcard = !this.fullCtx
	var merged = merge(existing.GetContext(), config.GetContext(), rootIsWildcard, mergeCache)
	// no need to check for existing.context, config.context in cache
	// since only way to create Newgraphs is "call rule" and here. We
	// cache at both places.
	existing.SetReachesIntoOuterContext(intMax(existing.GetReachesIntoOuterContext(), config.GetReachesIntoOuterContext()))
	// make sure to preserve the precedence filter suppression during the merge
	if config.getPrecedenceFilterSuppressed() {
		existing.setPrecedenceFilterSuppressed(true)
	}
	existing.SetContext(merged) // replace context no need to alt mapping
	return true
}

func (this *ATNConfigSet) GetStates() *Set {
	var states = NewSet(nil, nil)
	for i := 0; i < len(this.configs); i++ {
		states.add(this.configs[i].GetState())
	}
	return states
}

func (this *ATNConfigSet) getPredicates() []SemanticContext {
	var preds = make([]SemanticContext, 0)
	for i := 0; i < len(this.configs); i++ {
		c := this.configs[i].GetSemanticContext()
		if c != SemanticContextNONE {
			preds = append(preds, c)
		}
	}
	return preds
}

func (this *ATNConfigSet) getItems() []IATNConfig {
	return this.configs
}

func (this *ATNConfigSet) optimizeConfigs(interpreter *ATNSimulator) {
	if this.readOnly {
		panic("This set is readonly")
	}
	if this.configLookup.length() == 0 {
		return
	}
	for i := 0; i < len(this.configs); i++ {
		var config = this.configs[i]
		config.SetContext(interpreter.getCachedContext(config.GetContext()))
	}
}

func (this *ATNConfigSet) addAll(coll []*ATNConfig) bool {
	for i := 0; i < len(coll); i++ {
		this.add(coll[i], nil)
	}
	return false
}

func (this *ATNConfigSet) equals(other interface{}) bool {
	if this == other {
		return true
	} else if _, ok := other.(*ATNConfigSet); !ok {
		return false
	}

	other2 := other.(*ATNConfigSet)

	return this.configs != nil &&
		//			this.configs.equals(other2.configs) && // TODO is this necessary?
		this.fullCtx == other2.fullCtx &&
		this.uniqueAlt == other2.uniqueAlt &&
		this.conflictingAlts == other2.conflictingAlts &&
		this.hasSemanticContext == other2.hasSemanticContext &&
		this.dipsIntoOuterContext == other2.dipsIntoOuterContext
}

func (this *ATNConfigSet) hashString() string {
	if this.readOnly {
		if this.cachedHashString == "-1" {
			this.cachedHashString = this.hashConfigs()
		}
		return this.cachedHashString
	} else {
		return this.hashConfigs()
	}
}

func (this *ATNConfigSet) hashConfigs() string {
	var s = ""
	for _, c := range this.configs {
		s += fmt.Sprint(c)
	}
	return s
}

func (this *ATNConfigSet) length() int {
	return len(this.configs)
}

func (this *ATNConfigSet) isEmpty() bool {
	return len(this.configs) == 0
}

func (this *ATNConfigSet) contains(item *ATNConfig) bool {
	if this.configLookup == nil {
		panic("This method is not implemented for readonly sets.")
	}
	return this.configLookup.contains(item)
}

func (this *ATNConfigSet) containsFast(item *ATNConfig) bool {
	if this.configLookup == nil {
		panic("This method is not implemented for readonly sets.")
	}
	return this.configLookup.contains(item) // TODO containsFast is not implemented for Set
}

func (this *ATNConfigSet) clear() {
	if this.readOnly {
		panic("This set is readonly")
	}
	this.configs = make([]IATNConfig, 0)
	this.cachedHashString = "-1"
	this.configLookup = NewSet(hashATNConfig, equalATNConfigs)
}

func (this *ATNConfigSet) setReadonly(readOnly bool) {
	this.readOnly = readOnly
	if readOnly {
		this.configLookup = nil // can't mod, no need for lookup cache
	}
}

func (this *ATNConfigSet) toString() string {
	s := ""

	for _,c := range this.configs {
		s += c.toString()
	}

	if (this.hasSemanticContext){
		s += ",hasSemanticContext=" + fmt.Sprint(this.hasSemanticContext)
	}

	if (this.uniqueAlt != ATNINVALID_ALT_NUMBER ){
		s += ",uniqueAlt=" + fmt.Sprint(this.uniqueAlt)
	}

	if ( this.conflictingAlts != nil ){
		s += ",conflictingAlts=" + this.conflictingAlts.toString()
	}

	if (this.dipsIntoOuterContext) {
		s += ",dipsIntoOuterContext"
	}

	return s
}

type OrderedATNConfigSet struct {
	*ATNConfigSet
}

func NewOrderedATNConfigSet() *OrderedATNConfigSet {

	this := new(OrderedATNConfigSet)

	this.ATNConfigSet = NewATNConfigSet(false)
	this.configLookup = NewSet(nil, nil)

	return this
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/format"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"github.com/gedex/inflector"
)

var (
	needTimeImport bool
	structTypes    structTypeSlice
	outToStdout    = flag.Bool("c", false, "output to console; overrides \"-o\"")
	outputFile     = flag.String("o", "", "output file name; default is <schema>_schematype.go")
	packageName    = flag.String("p", "main", "package name for generated file; default is \"main\"")
)

type structField struct {
	Name         string
	Type         string
	Nullable     bool
	PropertyName string
	Required     bool
}

type structFieldSlice []structField

func (s structFieldSlice) Len() int {
	return len(s)
}

func (s structFieldSlice) Less(i, j int) bool {
	return s[i].Name < s[j].Name
}

func (s structFieldSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

type structType struct {
	Name    string
	Fields  structFieldSlice
	Comment string
}

type structTypeSlice []structType

func (s structTypeSlice) Len() int {
	return len(s)
}

func (s structTypeSlice) Less(i, j int) bool {
	return s[i].Name < s[j].Name
}

func (s structTypeSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func getTypeString(jsonType, format string) string {
	if format == "date-time" {
		needTimeImport = true
		return "time.Time"
	}

	switch jsonType {
	case "string":
		return "string"
	case "integer":
		return "int"
	case "number":
		return "float64"
	case "boolean":
		return "bool"
	case "null":
		return "nil"
	case "array":
		fallthrough
	case "object":
		return jsonType
	default:
		return "interface{}"
	}
}

// copied from golint (https://github.com/golang/lint/blob/4946cea8b6efd778dc31dc2dbeb919535e1b7529/lint.go#L701)
var commonInitialisms = map[string]bool{
	"API":   true,
	"ASCII": true,
	"CPU":   true,
	"CSS":   true,
	"DNS":   true,
	"EOF":   true,
	"GUID":  true,
	"HTML":  true,
	"HTTP":  true,
	"HTTPS": true,
	"ID":    true,
	"IP":    true,
	"JSON":  true,
	"LHS":   true,
	"QPS":   true,
	"RAM":   true,
	"RHS":   true,
	"RPC":   true,
	"SLA":   true,
	"SMTP":  true,
	"SQL":   true,
	"SSH":   true,
	"TCP":   true,
	"TLS":   true,
	"TTL":   true,
	"UDP":   true,
	"UI":    true,
	"UID":   true,
	"UUID":  true,
	"URI":   true,
	"URL":   true,
	"UTF8":  true,
	"VM":    true,
	"XML":   true,
	"XSRF":  true,
	"XSS":   true,
}

func dashedToWords(s string) string {
	return regexp.MustCompile("-|_").ReplaceAllString(s, " ")
}

func camelCaseToWords(s string) string {
	return regexp.MustCompile(`([\p{Ll}\p{N}])(\p{Lu})`).ReplaceAllString(s, "$1 $2")
}

func generateTypeName(origName string) string {
	spacedName := camelCaseToWords(dashedToWords(origName))
	titledName := strings.Title(spacedName)
	nameParts := strings.Split(titledName, " ")
	for i, part := range nameParts {
		upperedPart := strings.ToUpper(part)
		if commonInitialisms[upperedPart] {
			nameParts[i] = upperedPart
		}
	}
	return strings.Join(nameParts, "")
}

func getType(s *schema, pName, pDesc string) (typeName string) {
	var st structType

	if pName != "" {
		typeName = pName
	} else {
		typeName = s.Title
	}
	if st.Name = generateTypeName(typeName); st.Name == "" {
		log.Fatalln("Can't generate type without name.")
	}

	if pDesc != "" {
		st.Comment = pDesc
	} else {
		st.Comment = s.Description
	}

	required := make(map[string]bool)
	for _, req := range s.Required {
		required[req] = true
	}

	for propName, propSchema := range s.Properties {
		sf := structField{
			PropertyName: propName,
			Required:     required[propName],
		}

		var fieldName string
		if propSchema.Title != "" {
			fieldName = propSchema.Title
		} else {
			fieldName = propName
		}
		if sf.Name = generateTypeName(fieldName); sf.Name == "" {
			log.Fatalln("Can't generate field without name.")
		}

		switch propType := propSchema.Type.(type) {
		case []interface{}:
			if len(propType) == 2 && (propType[0] == "null" || propType[1] == "null") {
				sf.Nullable = true

				jsonType := propType[0]
				if jsonType == "null" {
					jsonType = propType[1]
				}

				sf.Type = getTypeString(jsonType.(string), propSchema.Format)
			}
		case string:
			if propType == "" {
				log.Fatalf("Can't create field %v withput type.", fieldName)
			}
			sf.Type = getTypeString(propType, propSchema.Format)
		case nil:
			log.Fatalf("Can't create field %v without type.", fieldName)
		}
		if sf.Type == "object" {
			sf.Type = getType(propSchema, sf.Name, propSchema.Description)
		} else if sf.Type == "array" {
			switch arrayItemType := propSchema.Items.(type) {
			case []interface{}:
				if len(arrayItemType) == 1 {
					singularName := inflector.Singularize(sf.Name)
					sf.Type = "[]" + getType(getArrayTypeSchema(arrayItemType[0]), singularName, propSchema.Description)
				} else {
					sf.Type = "[]interface{}"
				}
			case interface{}:
				singularName := inflector.Singularize(sf.Name)
				schema := getArrayTypeSchema(arrayItemType)
				sf.Type = "[]" + getType(schema, singularName, propSchema.Description)
			}
		}

		st.Fields = append(st.Fields, sf)
	}

	structTypes = append(structTypes, st)

	return
}

func getArrayTypeSchema(typeInterface interface{}) *schema {
	itemSchemaJSON, _ := json.Marshal(typeInterface)
	var itemSchema schema
	json.Unmarshal(itemSchemaJSON, &itemSchema)
	return &itemSchema
}

func (st structType) print(buf *bytes.Buffer) {
	if st.Comment != "" {
		buf.WriteString(fmt.Sprintln("//", st.Comment))
	}
	buf.WriteString(fmt.Sprintln("type ", st.Name, " struct {"))
	sort.Stable(st.Fields)
	for _, sf := range st.Fields {
		var typeString string
		if sf.Nullable && sf.Type != "interface{}" {
			typeString = "*"
		}
		typeString += sf.Type

		tagString := "`json:\"" + sf.PropertyName
		if !sf.Required {
			tagString += ",omitempty"
		}
		tagString += "\"`"
		buf.WriteString(fmt.Sprintln("   ", sf.Name, "  ", typeString, "  ", tagString))
	}
	buf.WriteString("}\n")
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		log.Fatalln("No file to parse.")
	}

	file, err := ioutil.ReadFile(flag.Arg(0))
	if err != nil {
		log.Fatalln("Error reading file:", err)
	}

	var s schema
	if err := json.Unmarshal(file, &s); err != nil {
		log.Fatalln("Error parsing JSON:", err)
	}

	schemaName := strings.Split(filepath.Base(flag.Arg(0)), ".")[0]
	getType(&s, schemaName, "")

	var resultSrc bytes.Buffer
	resultSrc.WriteString(fmt.Sprintln("package", *packageName))
	resultSrc.WriteString(fmt.Sprintf("\n// generated by \"%s\" -- DO NOT EDIT\n", strings.Join(os.Args, " ")))
	resultSrc.WriteString("\n")
	if needTimeImport {
		resultSrc.WriteString("import \"time\"\n")
	}
	sort.Stable(structTypes)
	for _, st := range structTypes {
		st.print(&resultSrc)
		resultSrc.WriteString("\n")
	}
	formattedSrc, err := format.Source(resultSrc.Bytes())
	if err != nil {
		log.Fatalln("Error running gofmt:", err)
	}

	if *outToStdout {
		fmt.Print(string(formattedSrc))
	} else {
		outputFileName := *outputFile
		if outputFileName == "" {
			compactSchemaName := strings.ToLower(generateTypeName(schemaName))
			outputFileName = fmt.Sprintf("%s_schematype.go", compactSchemaName)
		}
		err = ioutil.WriteFile(outputFileName, formattedSrc, 0644)
		if err != nil {
			log.Fatalf("Error writing to %s: %s\n", outputFileName, err)
		}
	}
}
<commit_msg>Handle objects without properties and untyped arrays<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/format"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"github.com/gedex/inflector"
)

var (
	needTimeImport bool
	structTypes    structTypeSlice
	outToStdout    = flag.Bool("c", false, "output to console; overrides \"-o\"")
	outputFile     = flag.String("o", "", "output file name; default is <schema>_schematype.go")
	packageName    = flag.String("p", "main", "package name for generated file; default is \"main\"")
)

type structField struct {
	Name         string
	Type         string
	Nullable     bool
	PropertyName string
	Required     bool
}

type structFieldSlice []structField

func (s structFieldSlice) Len() int {
	return len(s)
}

func (s structFieldSlice) Less(i, j int) bool {
	return s[i].Name < s[j].Name
}

func (s structFieldSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

type structType struct {
	Name    string
	Fields  structFieldSlice
	Comment string
}

type structTypeSlice []structType

func (s structTypeSlice) Len() int {
	return len(s)
}

func (s structTypeSlice) Less(i, j int) bool {
	return s[i].Name < s[j].Name
}

func (s structTypeSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func getTypeString(jsonType, format string) string {
	if format == "date-time" {
		needTimeImport = true
		return "time.Time"
	}

	switch jsonType {
	case "string":
		return "string"
	case "integer":
		return "int"
	case "number":
		return "float64"
	case "boolean":
		return "bool"
	case "null":
		return "nil"
	case "array":
		fallthrough
	case "object":
		return jsonType
	default:
		return "interface{}"
	}
}

// copied from golint (https://github.com/golang/lint/blob/4946cea8b6efd778dc31dc2dbeb919535e1b7529/lint.go#L701)
var commonInitialisms = map[string]bool{
	"API":   true,
	"ASCII": true,
	"CPU":   true,
	"CSS":   true,
	"DNS":   true,
	"EOF":   true,
	"GUID":  true,
	"HTML":  true,
	"HTTP":  true,
	"HTTPS": true,
	"ID":    true,
	"IP":    true,
	"JSON":  true,
	"LHS":   true,
	"QPS":   true,
	"RAM":   true,
	"RHS":   true,
	"RPC":   true,
	"SLA":   true,
	"SMTP":  true,
	"SQL":   true,
	"SSH":   true,
	"TCP":   true,
	"TLS":   true,
	"TTL":   true,
	"UDP":   true,
	"UI":    true,
	"UID":   true,
	"UUID":  true,
	"URI":   true,
	"URL":   true,
	"UTF8":  true,
	"VM":    true,
	"XML":   true,
	"XSRF":  true,
	"XSS":   true,
}

func dashedToWords(s string) string {
	return regexp.MustCompile("-|_").ReplaceAllString(s, " ")
}

func camelCaseToWords(s string) string {
	return regexp.MustCompile(`([\p{Ll}\p{N}])(\p{Lu})`).ReplaceAllString(s, "$1 $2")
}

func generateTypeName(origName string) string {
	spacedName := camelCaseToWords(dashedToWords(origName))
	titledName := strings.Title(spacedName)
	nameParts := strings.Split(titledName, " ")
	for i, part := range nameParts {
		upperedPart := strings.ToUpper(part)
		if commonInitialisms[upperedPart] {
			nameParts[i] = upperedPart
		}
	}
	return strings.Join(nameParts, "")
}

func getType(s *schema, pName, pDesc string) (typeName string) {
	var st structType

	if pName != "" {
		typeName = pName
	} else {
		typeName = s.Title
	}
	if st.Name = generateTypeName(typeName); st.Name == "" {
		log.Fatalln("Can't generate type without name.")
	}

	if pDesc != "" {
		st.Comment = pDesc
	} else {
		st.Comment = s.Description
	}

	required := make(map[string]bool)
	for _, req := range s.Required {
		required[req] = true
	}

	for propName, propSchema := range s.Properties {
		sf := structField{
			PropertyName: propName,
			Required:     required[propName],
		}

		var fieldName string
		if propSchema.Title != "" {
			fieldName = propSchema.Title
		} else {
			fieldName = propName
		}
		if sf.Name = generateTypeName(fieldName); sf.Name == "" {
			log.Fatalln("Can't generate field without name.")
		}

		switch propType := propSchema.Type.(type) {
		case []interface{}:
			if len(propType) == 2 && (propType[0] == "null" || propType[1] == "null") {
				sf.Nullable = true

				jsonType := propType[0]
				if jsonType == "null" {
					jsonType = propType[1]
				}

				sf.Type = getTypeString(jsonType.(string), propSchema.Format)
			}
		case string:
			if propType == "" {
				log.Fatalf("Can't create field %v withput type.", fieldName)
			}
			sf.Type = getTypeString(propType, propSchema.Format)
		case nil:
			log.Fatalf("Can't create field %v without type.", fieldName)
		}
		if sf.Type == "object" {
			if len(propSchema.Properties) > 0 {
				sf.Type = getType(propSchema, sf.Name, propSchema.Description)
			} else {
				sf.Type = "map[string]interface{}"
			}
		} else if sf.Type == "array" {
			switch arrayItemType := propSchema.Items.(type) {
			case []interface{}:
				if len(arrayItemType) == 1 {
					singularName := inflector.Singularize(sf.Name)
					sf.Type = "[]" + getType(getArrayTypeSchema(arrayItemType[0]), singularName, propSchema.Description)
				} else {
					sf.Type = "[]interface{}"
				}
			case interface{}:
				singularName := inflector.Singularize(sf.Name)
				schema := getArrayTypeSchema(arrayItemType)
				sf.Type = "[]" + getType(schema, singularName, propSchema.Description)
			default:
				sf.Type = "[]interface{}"
			}
		}

		st.Fields = append(st.Fields, sf)
	}

	structTypes = append(structTypes, st)

	return
}

func getArrayTypeSchema(typeInterface interface{}) *schema {
	itemSchemaJSON, _ := json.Marshal(typeInterface)
	var itemSchema schema
	json.Unmarshal(itemSchemaJSON, &itemSchema)
	return &itemSchema
}

func (st structType) print(buf *bytes.Buffer) {
	if st.Comment != "" {
		buf.WriteString(fmt.Sprintln("//", st.Comment))
	}
	buf.WriteString(fmt.Sprintln("type ", st.Name, " struct {"))
	sort.Stable(st.Fields)
	for _, sf := range st.Fields {
		var typeString string
		if sf.Nullable && sf.Type != "interface{}" {
			typeString = "*"
		}
		typeString += sf.Type

		tagString := "`json:\"" + sf.PropertyName
		if !sf.Required {
			tagString += ",omitempty"
		}
		tagString += "\"`"
		buf.WriteString(fmt.Sprintln("   ", sf.Name, "  ", typeString, "  ", tagString))
	}
	buf.WriteString("}\n")
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		log.Fatalln("No file to parse.")
	}

	file, err := ioutil.ReadFile(flag.Arg(0))
	if err != nil {
		log.Fatalln("Error reading file:", err)
	}

	var s schema
	if err := json.Unmarshal(file, &s); err != nil {
		log.Fatalln("Error parsing JSON:", err)
	}

	schemaName := strings.Split(filepath.Base(flag.Arg(0)), ".")[0]
	getType(&s, schemaName, "")

	var resultSrc bytes.Buffer
	resultSrc.WriteString(fmt.Sprintln("package", *packageName))
	resultSrc.WriteString(fmt.Sprintf("\n// generated by \"%s\" -- DO NOT EDIT\n", strings.Join(os.Args, " ")))
	resultSrc.WriteString("\n")
	if needTimeImport {
		resultSrc.WriteString("import \"time\"\n")
	}
	sort.Stable(structTypes)
	for _, st := range structTypes {
		st.print(&resultSrc)
		resultSrc.WriteString("\n")
	}
	formattedSrc, err := format.Source(resultSrc.Bytes())
	if err != nil {
		log.Fatalln("Error running gofmt:", err)
	}

	if *outToStdout {
		fmt.Print(string(formattedSrc))
	} else {
		outputFileName := *outputFile
		if outputFileName == "" {
			compactSchemaName := strings.ToLower(generateTypeName(schemaName))
			outputFileName = fmt.Sprintf("%s_schematype.go", compactSchemaName)
		}
		err = ioutil.WriteFile(outputFileName, formattedSrc, 0644)
		if err != nil {
			log.Fatalf("Error writing to %s: %s\n", outputFileName, err)
		}
	}
}
<|endoftext|>


<commit_before>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: 2020-09-11 03:26:24.428392 +0000 UTC

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "f4db670b9d9ffaa12bb9e1f7aca13a15833546b7"
<commit_msg>"Autogenerated build info"<commit_after>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: 2020-10-12 03:32:04.239376 +0000 UTC

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "7ba5c55c8c0b1109583e6e2a845090e1315e3067"
<|endoftext|>


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

package main

// TODO: -h, --human-readable
// TODO: -x, --one-file-system
// TODO: test symlinks

import (
	"fmt"
	"os"
	"sync"
	"sync/atomic"
	"syscall"

	"github.com/abiosoft/semaphore"
	flag "github.com/ogier/pflag"
)

// TODO: -k
var blockSize = flag.IntP("block-size", "B", 1024, "scale sizes by SIZE before printing them")
var summarize = flag.BoolP("summarize", "s", false, "display only a total for each argument")

var sem *semaphore.Semaphore

func init() {
	// TODO: size to number of fds
	// TODO: context instead of global
	sem = semaphore.New(512)
}

func Du(name ...string) (error, int64) {
	var wg sync.WaitGroup
	var sum int64
	var err error
	for _, arg := range name {
		wg.Add(1)
		go du(arg, &wg, &sum, &err)
	}
	wg.Wait()
	return err, sum
}

func du(name string, wgWg *sync.WaitGroup, sumSum *int64, err2 *error) error {
	defer wgWg.Done()

	var sum int64
	var wg sync.WaitGroup

	// avoid double release of semaphore
	semReleased := false
	semRelease := func() {
		if !semReleased {
			sem.Release()
			semReleased = true
		}
	}
	sem.Acquire()
	defer semRelease()

	dir, err := os.Open(name)
	if err != nil {
		// TODO: formatting different:
		// du: cannot access ‘/path’: Permission denied
		// du: lstat /path: permission denied
		fmt.Fprintf(os.Stderr, "du: %+v\n", err)
		sem.Release()
		// TODO: atomic set?
		*err2 = err
		return err
	}
	defer dir.Close()

	info, err := dir.Stat()
	if err != nil {
		fmt.Fprintf(os.Stderr, "du: %+v\n", err)
		*err2 = err
		return err
	}

	if stat, ok := info.Sys().(*syscall.Stat_t); ok {
		atomic.AddInt64(&sum, scaleBlocks(stat.Blocks))
	}

	if info.IsDir() {
		infos, err := dir.Readdir(0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "du: %+v\n", err)
			*err2 = err
			return err
		}
		dir.Close()
		semRelease()

		for _, info := range infos {
			fullname := name + "/" + info.Name()
			if info.IsDir() {
				wg.Add(1)
				go du(fullname, &wg, &sum, err2)
			} else if stat, ok := info.Sys().(*syscall.Stat_t); ok {
				atomic.AddInt64(&sum, scaleBlocks(stat.Blocks))
			}
		}
	}

	wg.Wait()
	atomic.AddInt64(sumSum, atomic.LoadInt64(&sum))
	if !*summarize {
		fmt.Printf("%d\t%s\n", atomic.LoadInt64(&sum), name)
	}
	return nil
}

func scaleBlocks(blocks int64) int64 {
	if *blockSize >= 512 {
		return blocks / (int64(*blockSize) / 512)
	} else {
		return blocks * (512 / int64(*blockSize))
	}
}

func main() {
	flag.Parse()
	args := flag.Args()
	if len(args) == 0 {
		args = append(args, ".")
	}

	err, sum := Du(args...)
	if *summarize {
		fmt.Printf("%d\t%s\n", sum, args[0])
	}
	if err != nil {
		os.Exit(1)
	}
}
<commit_msg>Add support for du --human-readable<commit_after>// Copyright 2017 Andrew Gaul
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

// TODO: -x, --one-file-system
// TODO: test symlinks

import (
	"fmt"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"syscall"

	"github.com/abiosoft/semaphore"
	flag "github.com/ogier/pflag"
)

// TODO: -k
var blockSize = flag.IntP("block-size", "B", 1024, "scale sizes by SIZE before printing them")
var humanReadable = flag.BoolP("human-readable", "h", false, "print sizes in human readable format (e.g., 1K 234M 2G)")
var summarize = flag.BoolP("summarize", "s", false, "display only a total for each argument")

var sem *semaphore.Semaphore

func init() {
	// TODO: size to number of fds
	// TODO: context instead of global
	sem = semaphore.New(512)
}

func Du(name ...string) (error, int64) {
	var wg sync.WaitGroup
	var sum int64
	var err error
	for _, arg := range name {
		wg.Add(1)
		go du(arg, &wg, &sum, &err)
	}
	wg.Wait()
	return err, sum
}

func du(name string, wgWg *sync.WaitGroup, sumSum *int64, err2 *error) error {
	defer wgWg.Done()

	var sum int64
	var wg sync.WaitGroup

	// avoid double release of semaphore
	semReleased := false
	semRelease := func() {
		if !semReleased {
			sem.Release()
			semReleased = true
		}
	}
	sem.Acquire()
	defer semRelease()

	dir, err := os.Open(name)
	if err != nil {
		// TODO: formatting different:
		// du: cannot access ‘/path’: Permission denied
		// du: lstat /path: permission denied
		fmt.Fprintf(os.Stderr, "du: %+v\n", err)
		sem.Release()
		// TODO: atomic set?
		*err2 = err
		return err
	}
	defer dir.Close()

	info, err := dir.Stat()
	if err != nil {
		fmt.Fprintf(os.Stderr, "du: %+v\n", err)
		*err2 = err
		return err
	}

	if stat, ok := info.Sys().(*syscall.Stat_t); ok {
		atomic.AddInt64(&sum, scaleBlocks(stat.Blocks))
	}

	if info.IsDir() {
		infos, err := dir.Readdir(0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "du: %+v\n", err)
			*err2 = err
			return err
		}
		dir.Close()
		semRelease()

		for _, info := range infos {
			fullname := name + "/" + info.Name()
			if info.IsDir() {
				wg.Add(1)
				go du(fullname, &wg, &sum, err2)
			} else if stat, ok := info.Sys().(*syscall.Stat_t); ok {
				atomic.AddInt64(&sum, scaleBlocks(stat.Blocks))
			}
		}
	}

	wg.Wait()
	atomic.AddInt64(sumSum, atomic.LoadInt64(&sum))
	if !*summarize {
		fmt.Printf("%s\t%s\n", formatBlocks(atomic.LoadInt64(&sum)), name)
	}
	return nil
}

func scaleBlocks(blocks int64) int64 {
	if *blockSize >= 512 {
		return blocks / (int64(*blockSize) / 512)
	} else {
		return blocks * (512 / int64(*blockSize))
	}
}

func formatBlocks(blocks int64) string {
	if *humanReadable {
		bytes := blocks * int64(*blockSize)
		if bytes > (1 << 60) {
			return fmt.Sprintf("%2.1fE", float64(bytes)/(1<<60))
		} else if bytes > 10*(1<<50) {
			return fmt.Sprintf("%dP", bytes/(1<<50))
		} else if bytes > (1 << 50) {
			return fmt.Sprintf("%2.1fP", float64(bytes)/(1<<50))
		} else if bytes > 10*(1<<40) {
			return fmt.Sprintf("%dT", bytes/(1<<40))
		} else if bytes > (1 << 40) {
			return fmt.Sprintf("%2.1fT", float64(bytes)/(1<<40))
		} else if bytes > 10*(1<<30) {
			return fmt.Sprintf("%dG", bytes/(1<<30))
		} else if bytes > (1 << 30) {
			return fmt.Sprintf("%2.1fG", float64(bytes)/(1<<30))
		} else if bytes > 10*(1<<20) {
			return fmt.Sprintf("%dM", bytes/(1<<20))
		} else if bytes > (1 << 20) {
			return fmt.Sprintf("%2.1fM", float64(bytes)/(1<<20))
		} else if bytes > 10*(1<<10) {
			return fmt.Sprintf("%dK", bytes/(1<<10))
		} else if bytes > (1 << 10) {
			return fmt.Sprintf("%2.1fK", float64(bytes)/(1<<10))
		} else {
			return strconv.FormatInt(bytes, 10)
		}
	} else {
		return strconv.FormatInt(blocks, 10)
	}
}

func main() {
	flag.Parse()
	args := flag.Args()
	if len(args) == 0 {
		args = append(args, ".")
	}

	err, sum := Du(args...)
	if *summarize {
		fmt.Printf("%s\t%s\n", formatBlocks(sum), args[0])
	}
	if err != nil {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before><commit_msg>changes to inject of otelhttp transport in executor (#2552)<commit_after><|endoftext|>