<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 exec

import (
	"context"
	"fmt"
	"io"
	"math"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
)

// DataSource is a Root execution unit.
type DataSource struct {
	UID   UnitID
	SID   StreamID
	Coder *coder.Coder
	Out   Node

	source   DataManager
	count    int64
	splitPos int64
	start    time.Time

	mu sync.Mutex
}

func (n *DataSource) ID() UnitID {
	return n.UID
}

func (n *DataSource) Up(ctx context.Context) error {
	return nil
}

func (n *DataSource) StartBundle(ctx context.Context, id string, data DataContext) error {
	n.mu.Lock()
	n.source = data.Data
	n.start = time.Now()
	n.count = 0
	n.splitPos = math.MaxInt64
	n.mu.Unlock()
	return n.Out.StartBundle(ctx, id, data)
}

func (n *DataSource) Process(ctx context.Context) error {
	r, err := n.source.OpenRead(ctx, n.SID)
	if err != nil {
		return err
	}
	defer r.Close()

	c := coder.SkipW(n.Coder)
	wc := MakeWindowDecoder(n.Coder.Window)

	switch {
	case coder.IsCoGBK(c):
		ck := MakeElementDecoder(c.Components[0])
		cv := MakeElementDecoder(c.Components[1])

		for {
			if n.IncrementCountAndCheckSplit(ctx) {
				return nil
			}
			ws, t, err := DecodeWindowedValueHeader(wc, r)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return errors.Wrap(err, "source failed")
			}

			// Decode key

			key, err := ck.Decode(r)
			if err != nil {
				return errors.Wrap(err, "source decode failed")
			}
			key.Timestamp = t
			key.Windows = ws

			// TODO(herohde) 4/30/2017: the State API will be handle re-iterations
			// and only "small" value streams would be inline. Presumably, that
			// would entail buffering the whole stream. We do that for now.

			var buf []FullValue

			size, err := coder.DecodeInt32(r)
			if err != nil {
				return errors.Wrap(err, "stream size decoding failed")
			}

			if size > -1 {
				// Single chunk stream.

				// log.Printf("Fixed size=%v", size)
				for i := int32(0); i < size; i++ {
					value, err := cv.Decode(r)
					if err != nil {
						return errors.Wrap(err, "stream value decode failed")
					}
					buf = append(buf, *value)
				}
			} else {
				// Multi-chunked stream.

				for {
					chunk, err := coder.DecodeVarUint64(r)
					if err != nil {
						return errors.Wrap(err, "stream chunk size decoding failed")
					}

					// log.Printf("Chunk size=%v", chunk)

					if chunk == 0 {
						break
					}

					for i := uint64(0); i < chunk; i++ {
						value, err := cv.Decode(r)
						if err != nil {
							return errors.Wrap(err, "stream value decode failed")
						}
						buf = append(buf, *value)
					}
				}
			}

			values := &FixedReStream{Buf: buf}
			if err := n.Out.ProcessElement(ctx, key, values); err != nil {
				return err
			}
		}

	default:
		ec := MakeElementDecoder(c)

		for {
			if n.IncrementCountAndCheckSplit(ctx) {
				return nil
			}

			ws, t, err := DecodeWindowedValueHeader(wc, r)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return errors.Wrap(err, "source failed")
			}

			elm, err := ec.Decode(r)
			if err != nil {
				return errors.Wrap(err, "source decode failed")
			}
			elm.Timestamp = t
			elm.Windows = ws

			// log.Printf("READ: %v %v", elm.Key.Type(), elm.Key.Interface())

			if err := n.Out.ProcessElement(ctx, elm); err != nil {
				return err
			}
		}
	}
}

func (n *DataSource) FinishBundle(ctx context.Context) error {
	n.mu.Lock()
	log.Infof(ctx, "DataSource: %d elements in %d ns", n.count, time.Now().Sub(n.start))
	n.source = nil
	err := n.Out.FinishBundle(ctx)
	n.count = 0
	n.splitPos = math.MaxInt64
	defer n.mu.Unlock()
	return err
}

func (n *DataSource) Down(ctx context.Context) error {
	n.source = nil
	return nil
}

func (n *DataSource) String() string {
	return fmt.Sprintf("DataSource[%v] Coder:%v Out:%v", n.SID, n.Coder, n.Out.ID())
}

// IncrementCountAndCheckSplit increments DataSource.count by one and checks if
// the caller should abort further element processing, and finish the bundle.
// Returns true if the new value of count is greater than or equal to the split
// point, and false otherwise.
func (n *DataSource) IncrementCountAndCheckSplit(ctx context.Context) bool {
	b := false
	n.mu.Lock()
	n.count++
	if n.count >= n.splitPos {
		b = true
	}
	n.mu.Unlock()
	return b
}

// ProgressReportSnapshot captures the progress reading an input source.
type ProgressReportSnapshot struct {
	ID, Name string
	Count    int64
}

// Progress returns a snapshot of the source's progress.
func (n *DataSource) Progress() ProgressReportSnapshot {
	if n == nil {
		return ProgressReportSnapshot{}
	}
	n.mu.Lock()
	c := n.count
	n.mu.Unlock()
	return ProgressReportSnapshot{n.SID.Target.ID, n.SID.Target.Name, c}
}

// Split takes a sorted set of potential split points, selects and actuates
// split on an appropriate split point, and returns the selected split point
// if successful. Returns an error when unable to split.
func (n *DataSource) Split(splits []int64, frac float32) (int64, error) {
	if splits == nil {
		return 0, fmt.Errorf("failed to split: requested splits were empty")
	}
	if n == nil {
		return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource not initialized", splits)
	}
	n.mu.Lock()
	c := n.count
	// Find the smallest split index that we haven't yet processed, and set
	// the promised split position to this value.
	for _, s := range splits {
		if s >= c {
			n.splitPos = s
			fs := n.splitPos
			n.mu.Unlock()
			return fs, nil
		}
	}
	n.mu.Unlock()
	// If we can't find a suitable split point from the requested choices,
	// return an error.
	return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource at index: %v", splits, c)
}
<commit_msg>Update datasource.go<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 exec

import (
	"context"
	"fmt"
	"io"
	"math"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/graph/coder"
	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
)

// DataSource is a Root execution unit.
type DataSource struct {
	UID   UnitID
	SID   StreamID
	Coder *coder.Coder
	Out   Node

	source   DataManager
	count    int64
	splitPos int64
	start    time.Time

	mu sync.Mutex
}

func (n *DataSource) ID() UnitID {
	return n.UID
}

func (n *DataSource) Up(ctx context.Context) error {
	return nil
}

func (n *DataSource) StartBundle(ctx context.Context, id string, data DataContext) error {
	n.mu.Lock()
	n.source = data.Data
	n.start = time.Now()
	n.count = 0
	n.splitPos = math.MaxInt64
	n.mu.Unlock()
	return n.Out.StartBundle(ctx, id, data)
}

func (n *DataSource) Process(ctx context.Context) error {
	r, err := n.source.OpenRead(ctx, n.SID)
	if err != nil {
		return err
	}
	defer r.Close()

	c := coder.SkipW(n.Coder)
	wc := MakeWindowDecoder(n.Coder.Window)

	switch {
	case coder.IsCoGBK(c):
		ck := MakeElementDecoder(c.Components[0])
		cv := MakeElementDecoder(c.Components[1])

		for {
			if n.IncrementCountAndCheckSplit(ctx) {
				return nil
			}
			ws, t, err := DecodeWindowedValueHeader(wc, r)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return errors.Wrap(err, "source failed")
			}

			// Decode key

			key, err := ck.Decode(r)
			if err != nil {
				return errors.Wrap(err, "source decode failed")
			}
			key.Timestamp = t
			key.Windows = ws

			// TODO(herohde) 4/30/2017: the State API will be handle re-iterations
			// and only "small" value streams would be inline. Presumably, that
			// would entail buffering the whole stream. We do that for now.

			var buf []FullValue

			size, err := coder.DecodeInt32(r)
			if err != nil {
				return errors.Wrap(err, "stream size decoding failed")
			}

			if size > -1 {
				// Single chunk stream.

				// log.Printf("Fixed size=%v", size)
				for i := int32(0); i < size; i++ {
					value, err := cv.Decode(r)
					if err != nil {
						return errors.Wrap(err, "stream value decode failed")
					}
					buf = append(buf, *value)
				}
			} else {
				// Multi-chunked stream.

				for {
					chunk, err := coder.DecodeVarUint64(r)
					if err != nil {
						return errors.Wrap(err, "stream chunk size decoding failed")
					}

					// log.Printf("Chunk size=%v", chunk)

					if chunk == 0 {
						break
					}

					for i := uint64(0); i < chunk; i++ {
						value, err := cv.Decode(r)
						if err != nil {
							return errors.Wrap(err, "stream value decode failed")
						}
						buf = append(buf, *value)
					}
				}
			}

			values := &FixedReStream{Buf: buf}
			if err := n.Out.ProcessElement(ctx, key, values); err != nil {
				return err
			}
		}

	default:
		ec := MakeElementDecoder(c)

		for {
			if n.IncrementCountAndCheckSplit(ctx) {
				return nil
			}

			ws, t, err := DecodeWindowedValueHeader(wc, r)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return errors.Wrap(err, "source failed")
			}

			elm, err := ec.Decode(r)
			if err != nil {
				return errors.Wrap(err, "source decode failed")
			}
			elm.Timestamp = t
			elm.Windows = ws

			// log.Printf("READ: %v %v", elm.Key.Type(), elm.Key.Interface())

			if err := n.Out.ProcessElement(ctx, elm); err != nil {
				return err
			}
		}
	}
}

func (n *DataSource) FinishBundle(ctx context.Context) error {
	n.mu.Lock()
	defer n.mu.Unlock()
	log.Infof(ctx, "DataSource: %d elements in %d ns", n.count, time.Now().Sub(n.start))
	n.source = nil
	err := n.Out.FinishBundle(ctx)
	n.count = 0
	n.splitPos = math.MaxInt64
	return err
}

func (n *DataSource) Down(ctx context.Context) error {
	n.source = nil
	return nil
}

func (n *DataSource) String() string {
	return fmt.Sprintf("DataSource[%v] Coder:%v Out:%v", n.SID, n.Coder, n.Out.ID())
}

// IncrementCountAndCheckSplit increments DataSource.count by one and checks if
// the caller should abort further element processing, and finish the bundle.
// Returns true if the new value of count is greater than or equal to the split
// point, and false otherwise.
func (n *DataSource) IncrementCountAndCheckSplit(ctx context.Context) bool {
	b := false
	n.mu.Lock()
	n.count++
	if n.count >= n.splitPos {
		b = true
	}
	n.mu.Unlock()
	return b
}

// ProgressReportSnapshot captures the progress reading an input source.
type ProgressReportSnapshot struct {
	ID, Name string
	Count    int64
}

// Progress returns a snapshot of the source's progress.
func (n *DataSource) Progress() ProgressReportSnapshot {
	if n == nil {
		return ProgressReportSnapshot{}
	}
	n.mu.Lock()
	c := n.count
	n.mu.Unlock()
	return ProgressReportSnapshot{n.SID.Target.ID, n.SID.Target.Name, c}
}

// Split takes a sorted set of potential split points, selects and actuates
// split on an appropriate split point, and returns the selected split point
// if successful. Returns an error when unable to split.
func (n *DataSource) Split(splits []int64, frac float32) (int64, error) {
	if splits == nil {
		return 0, fmt.Errorf("failed to split: requested splits were empty")
	}
	if n == nil {
		return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource not initialized", splits)
	}
	n.mu.Lock()
	c := n.count
	// Find the smallest split index that we haven't yet processed, and set
	// the promised split position to this value.
	for _, s := range splits {
		if s >= c {
			n.splitPos = s
			fs := n.splitPos
			n.mu.Unlock()
			return fs, nil
		}
	}
	n.mu.Unlock()
	// If we can't find a suitable split point from the requested choices,
	// return an error.
	return 0, fmt.Errorf("failed to split at requested splits: {%v}, DataSource at index: %v", splits, c)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"
)

var hexRegex = regexp.MustCompile(`^[a-f\d]{6}$`)

var colors = map[string]string{
	"bg-none":    "0",
	"fg-black":   "30",
	"fg-red":     "31",
	"fg-green":   "32",
	"fg-yellow":  "33",
	"fg-blue":    "34",
	"fg-magenta": "35",
	"fg-cyan":    "36",
	"fg-white":   "37",
	"bg-black":   "40",
	"bg-red":     "41",
	"bg-green":   "42",
	"bg-yellow":  "43",
	"bg-blue":    "44",
	"bg-magenta": "45",
	"bg-cyan":    "46",
	"bg-white":   "47",
}

func escapeBackground(color string) string {
	switch shell {
	case "zsh":
		return "%K{" + color + "}"
	case "bash":
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[48;2;" + escapeHex(color) + "m"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func escapeForeground(color string) string {
	switch shell {
	case "zsh":
		return "%F{" + color + "}"
	case "bash":
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[38;2;" + escapeHex(color) + "m"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func resetColors() {
	switch shell {
	case "zsh":
		fmt.Print("%{%f%}")
	case "bash":
		fmt.Print("\\[\x1b[0m\\]")
	default:
		fmt.Print("\x1b[0m")
	}
}

func escapeHex(color string) string {
	var r, g, b byte
	fmt.Sscanf(color, "%02x%02x%02x", &r, &g, &b)
	return fmt.Sprintf("%d;%d;%d", r, g, b)
}
<commit_msg>added 24-bit color support for zsh<commit_after>package main

import (
	"fmt"
	"regexp"
)

var hexRegex = regexp.MustCompile(`^[a-f\d]{6}$`)

var colors = map[string]string{
	"bg-none":    "0",
	"fg-black":   "30",
	"fg-red":     "31",
	"fg-green":   "32",
	"fg-yellow":  "33",
	"fg-blue":    "34",
	"fg-magenta": "35",
	"fg-cyan":    "36",
	"fg-white":   "37",
	"bg-black":   "40",
	"bg-red":     "41",
	"bg-green":   "42",
	"bg-yellow":  "43",
	"bg-blue":    "44",
	"bg-magenta": "45",
	"bg-cyan":    "46",
	"bg-white":   "47",
}

func escapeBackground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[48;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%K{" + color + "}"
	case "bash":
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[48;2;" + escapeHex(color) + "m"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func escapeForeground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[38;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%F{" + color + "}"
	case "bash":
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[38;2;" + escapeHex(color) + "m"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func resetColors() {
	switch shell {
	case "zsh":
		fmt.Print("%{%f%}")
	case "bash":
		fmt.Print("\\[\x1b[0m\\]")
	default:
		fmt.Print("\x1b[0m")
	}
}

func escapeHex(color string) string {
	var r, g, b byte
	fmt.Sscanf(color, "%02x%02x%02x", &r, &g, &b)
	return fmt.Sprintf("%d;%d;%d", r, g, b)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"
)

var hexRegex = regexp.MustCompile(`^[a-f\d]{6}$`)

var colors = map[string]string{
	"bg-none":    "0",
	"fg-black":   "30",
	"fg-red":     "31",
	"fg-green":   "32",
	"fg-yellow":  "33",
	"fg-blue":    "34",
	"fg-magenta": "35",
	"fg-cyan":    "36",
	"fg-white":   "37",
	"bg-black":   "40",
	"bg-red":     "41",
	"bg-green":   "42",
	"bg-yellow":  "43",
	"bg-blue":    "44",
	"bg-magenta": "45",
	"bg-cyan":    "46",
	"bg-white":   "47",
}

func escapeBackground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[48;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%K{" + color + "}"
	case "bash":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\\[\x1b[48;2;" + escapeHex(color) + "m\\]"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[48;2;" + escapeHex(color) + "m"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func escapeForeground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[38;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%F{" + color + "}"
	case "bash":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\\[\x1b[38;2;" + escapeHex(color) + "m\\]"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[38;2;" + escapeHex(color) + "m"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func resetColors() {
	switch shell {
	case "zsh":
		fmt.Print("%{%f%}")
	case "bash":
		fmt.Print("\\[\x1b[0m\\]")
	default:
		fmt.Print("\x1b[0m")
	}
}

func escapeHex(color string) string {
	var r, g, b byte
	fmt.Sscanf(color, "%02x%02x%02x", &r, &g, &b)
	return fmt.Sprintf("%d;%d;%d", r, g, b)
}
<commit_msg>added support for 256 colors<commit_after>package main

import (
	"fmt"
	"regexp"
	"strconv"
)

var hexRegex = regexp.MustCompile(`^[a-f\d]{6}$`)

var colors = map[string]string{
	"bg-none":    "0",
	"fg-black":   "30",
	"fg-red":     "31",
	"fg-green":   "32",
	"fg-yellow":  "33",
	"fg-blue":    "34",
	"fg-magenta": "35",
	"fg-cyan":    "36",
	"fg-white":   "37",
	"bg-black":   "40",
	"bg-red":     "41",
	"bg-green":   "42",
	"bg-yellow":  "43",
	"bg-blue":    "44",
	"bg-magenta": "45",
	"bg-cyan":    "46",
	"bg-white":   "47",
}

func escapeBackground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[48;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%K{" + color + "}"
	case "bash":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\\[\x1b[48;2;" + escapeHex(color) + "m\\]"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[48;2;" + escapeHex(color) + "m"
		}

		// 256 colors
		_, err := strconv.Atoi(color)
		if err == nil {
			return "\x1b[48;5;" + color + "m"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func escapeForeground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[38;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%F{" + color + "}"
	case "bash":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\\[\x1b[38;2;" + escapeHex(color) + "m\\]"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[38;2;" + escapeHex(color) + "m"
		}

		// 256 colors
		_, err := strconv.Atoi(color)
		if err == nil {
			return "\x1b[38;5;" + color + "m"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func resetColors() {
	switch shell {
	case "zsh":
		fmt.Print("%{%f%}")
	case "bash":
		fmt.Print("\\[\x1b[0m\\]")
	default:
		fmt.Print("\x1b[0m")
	}
}

func escapeHex(color string) string {
	var r, g, b byte
	fmt.Sscanf(color, "%02x%02x%02x", &r, &g, &b)
	return fmt.Sprintf("%d;%d;%d", r, g, b)
}
<|endoftext|>


<commit_before>package kbmap

import (
	//"fmt"
	"github.com/driusan/de/actions"
	"github.com/driusan/de/demodel"
	"github.com/driusan/de/position"
	"golang.org/x/mobile/event/key"
)

var Repeat uint

func normalMap(e key.Event, buff *demodel.CharBuffer, v demodel.Viewport) (demodel.Map, demodel.ScrollDirection, error) {
	// things only happen on key press in normal mode, if it's a release
	// or a repeat, ignore it. It's not an error
	if e.Direction == key.DirRelease {
		return NormalMode, demodel.DirectionNone, nil
	}
	if buff == nil {
		return NormalMode, demodel.DirectionNone, ErrInvalid
	}
	switch e.Code {
	case key.CodeEscape:
		actions.Do("SaveOrExit", buff, v)
		return NormalMode, demodel.DirectionUp, nil
	case key.CodeDeleteBackspace:
		if e.Direction == key.DirPress {
			actions.DeleteCursor(position.DotStart, position.DotEnd, buff)
		}
		// There's a potential that we're pressing backspace at the start of the
		// screen and may need to scroll up
		return NormalMode, demodel.DirectionUp, nil
	case key.CodeI:
		return InsertMode, demodel.DirectionNone, nil
	case key.CodeA:
		if buff.Dot.End >= uint(len(buff.Buffer)) {
			return InsertMode, demodel.DirectionDown, nil
		}
		if buff.Buffer[buff.Dot.End] == '\n' {
			// vi doesn't actually let you navigate your cursor *on top*
			// of the newline, while de does.
			// It feels weird to push 'a' there and end up on the next
			// line, so if we append to the end of a line, just leave
			// the cursor end where it is,
			actions.MoveCursor(position.DotEnd, position.DotEnd, buff)
		} else {
			actions.MoveCursor(position.NextChar, position.NextChar, buff)

		}

		// There's a potentially we pressed 'a' at the very end of the screen and
		// need to scroll down
		return InsertMode, demodel.DirectionDown, nil
	case key.CodeK:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				// ctrl is pressed, so just move the start and expand dot
				actions.MoveCursor(position.PrevLine, position.DotEnd, buff)
			} else {
				// ctrl is not pressed, so move the cursor without selecting
				actions.MoveCursor(position.PrevLine, position.PrevLine, buff)

			}
		}

		return NormalMode, demodel.DirectionUp, nil
	case key.CodeH:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				// ctrl is pressed, so just move the end and expand dot
				actions.MoveCursor(position.PrevChar, position.DotEnd, buff)
			} else {
				// ctrl is not pressed, so move the cursor without selecting
				actions.MoveCursor(position.PrevChar, position.PrevChar, buff)

			}
		}

		return NormalMode, demodel.DirectionUp, nil

	case key.CodeL:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				// ctrl is pressed, so just move the end and expand dot
				actions.MoveCursor(position.DotStart, position.NextChar, buff)
			} else {
				// ctrl is not pressed, so move the cursor without selecting
				actions.MoveCursor(position.NextChar, position.NextChar, buff)
			}
		}

		return NormalMode, demodel.DirectionDown, nil
	case key.CodeJ:
		if e.Modifiers&key.ModShift != 0 {
			// 'J' for Join

			undoDot := -1
			if buff.Dot.Start == buff.Dot.End {
				if Repeat <= 1 {
					Repeat = 2
				}
				undoDot = int(buff.Dot.Start)
				// Join n lines
				actions.MoveCursor(position.StartOfLine, position.DotEnd, buff)
				for ; Repeat > 1; Repeat-- {
					actions.MoveCursor(position.DotStart, position.NextLine, buff)
				}
				actions.MoveCursor(position.DotStart, position.EndOfLine, buff)
			}

			buff.JoinLines(buff.Dot.Start, buff.Dot.End)
			if undoDot >= 0 {
				buff.Undo.Dot.Start = uint(undoDot)
				buff.Undo.Dot.End = uint(undoDot)
				buff.Dot = buff.Undo.Dot
			}
			v.InvalidateCache()
			Repeat = 0
			return NormalMode, demodel.DirectionNone, nil
		}

		// 'j' for movement
		if Repeat == 0 {
			Repeat = 1

		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				actions.MoveCursor(position.DotStart, position.NextLine, buff)
			} else {
				actions.MoveCursor(position.NextLine, position.NextLine, buff)
			}
		}

		return NormalMode, demodel.DirectionDown, nil
	case key.CodeX:
		if e.Direction == key.DirPress && isCopyModifier(e) {
			actions.CutClipboard(position.DotStart, position.DotEnd, buff)
			return NormalMode, demodel.DirectionNone, nil
		}

		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			actions.MoveCursor(position.DotStart, position.NextChar, buff)
		}
		actions.DeleteCursor(position.DotStart, position.DotEnd, buff)
		return NormalMode, demodel.DirectionNone, nil
	case key.CodeW:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				actions.MoveCursor(position.DotStart, position.NextWordStart, buff)
			} else {
				actions.MoveCursor(position.NextWordStart, position.NextWordStart, buff)
			}
		}
		return NormalMode, demodel.DirectionDown, nil
	case key.CodeB:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				actions.MoveCursor(position.PrevWordStart, position.DotEnd, buff)
			} else {
				actions.MoveCursor(position.PrevWordStart, position.PrevWordStart, buff)
			}
		}
		return NormalMode, demodel.DirectionUp, nil
	case key.CodeP:
		actions.InsertSnarfBuffer(position.DotStart, position.DotEnd, buff)
		return NormalMode, demodel.DirectionDown, nil
	case key.CodeRightArrow:
		// Arrow keys indicate their scroll direction via the error return value,
		// they return demodel.DirectionNone to make sure both code paths don't accidentally
		// get triggered
		return NormalMode, demodel.DirectionNone, ErrScrollRight
	case key.CodeLeftArrow:
		return NormalMode, demodel.DirectionNone, ErrScrollLeft
	case key.CodeDownArrow:
		return NormalMode, demodel.DirectionNone, ErrScrollDown
	case key.CodeUpArrow:
		return NormalMode, demodel.DirectionNone, ErrScrollUp
	case key.Code0:
		if e.Modifiers == 0 {
			Repeat *= 10
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code1:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 1
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code2:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 2

		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code3:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 3

		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code4:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 4

		} else if e.Modifiers&key.ModShift != 0 {
			// $ is pressed, in most key maps..
			// it doesn't mean anything to repeat "End of line", so just
			// reset the repeat counter.
			Repeat = 0

			if e.Modifiers&key.ModControl != 0 {
				// ctrl is pressed, so just move the end and expand dot
				// BUG: This doesn't seem to work. Probably a bug in the shiny driver.
				actions.MoveCursor(position.DotStart, position.EndOfLine, buff)
			} else {
				actions.MoveCursor(position.EndOfLine, position.EndOfLine, buff)
			}

		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code5:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 5
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code6:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 6
		} else if e.Modifiers&key.ModShift != 0 {
			// ^ is pressed, on most keyboards..
			// It doesn't mean anything to repeat "Start of line" either
			Repeat = 0

			if e.Modifiers&key.ModControl != 0 {
				// BUG: This doesn't seem to work. Probably a bug in the shiny driver.
				actions.MoveCursor(position.StartOfLine, position.DotEnd, buff)
			} else {
				actions.MoveCursor(position.StartOfLine, position.StartOfLine, buff)
			}

		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code7:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 7
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code8:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 8
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code9:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 9
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.CodeG:
		if e.Modifiers&key.ModShift != 0 {
			// G treats the repeat counter differently. If it's set, it means "Go to
			// lineno", and if it's not it means "Go to end of buffer"
			if Repeat > 0 {
				/* This seems to be buggy, so do it manually, at least until there's better
				   tests.
				actions.MoveCursor(position.BuffStart, position.BuffStart, buff)
				for ; Repeat > 0; Repeat-- {
					actions.MoveCursor(position.NextLine, position.NextLine, buff)
				}*/
				lineNo := uint(1)
				for i, c := range buff.Buffer {
					if c == '\n' {
						lineNo++
					}
					// should be equal, but if we've went too far at least it's
					// better to stop now.
					if lineNo >= Repeat {
						// we're on the \n itself, so add 1 to get to the actual line.
						if Repeat == 1 {
							buff.Dot.Start = 0
						} else {
							buff.Dot.Start = uint(i + 1)
						}
						buff.Dot.End = buff.Dot.Start
						Repeat = 0
						// select the line.
						actions.MoveCursor(position.DotStart, position.EndOfLine, buff)
						break
					}
				}
				// We don't know if the line being moved to is above or below, but send an "up"
				// hint so that it centers it at the top of the screen based on dot.
				return NormalMode, demodel.DirectionUp, nil
			}
			// moving to the end of the file because Repeat isn't set.
			actions.MoveCursor(position.BuffEnd, position.BuffEnd, buff)
			// send a down hint so it centers it at the bottom.
			return NormalMode, demodel.DirectionDown, nil
		}

	case key.CodeD:
		return DeleteMode, demodel.DirectionNone, nil
	case key.CodeU:
		actions.Do("Undo", buff, v)
		return NormalMode, demodel.DirectionNone, nil
	case key.CodeO:
		actions.MoveCursor(position.EndOfLineWithNewline, position.EndOfLineWithNewline, buff)

		newBuffer := make([]byte, len(buff.Buffer)+1)
		copy(newBuffer, buff.Buffer[:buff.Dot.Start])
		newBuffer[buff.Dot.Start] = '\n'
		copy(newBuffer[buff.Dot.Start+1:], buff.Buffer[buff.Dot.End:])

		buff.Buffer = newBuffer
		return InsertMode, demodel.DirectionDown, nil
	case key.CodeSlash:
		if buff.Dot.Start == buff.Dot.End {
			actions.FindNext(position.CurWordStart, position.CurWordEnd, buff)

		} else {
			actions.FindNext(position.DotStart, position.DotEnd, buff)
		}
		return NormalMode, demodel.DirectionDown, nil
	case key.CodeReturnEnter:
		if buff.Dot.Start == buff.Dot.End {
			actions.OpenOrPerformAction(position.CurExecutionWordStart, position.CurExecutionWordEnd, buff, v)
		} else {
			actions.OpenOrPerformAction(position.DotStart, position.DotEnd, buff, v)
		}
		// There's a possibility OpenOrPerformAction opened a new file, in which case
		// we should scroll to the top, or inserted text, in which case we should scroll
		// to the top of the inserted text. Either way, it's an "Up" hint so that the
		// scrolling is based on Dot.Start
		return v.GetKeyboardMode(), demodel.DirectionUp, nil
	case key.CodeSemicolon:
		if e.Modifiers&key.ModShift != 0 {
			if buff.Tagline != nil {
				buff.Tagline.Buffer = append(buff.Tagline.Buffer, ' ')

				buff.Tagline.Dot.Start = uint(len(buff.Tagline.Buffer))
				buff.Tagline.Dot.End = buff.Tagline.Dot.Start
			}
		}
		return TagMode, demodel.DirectionNone, nil
	case key.CodeV:
		if e.Direction == key.DirPress && isCopyModifier(e) {
			actions.PasteClipboard(position.DotStart, position.DotEnd, buff)
			return NormalMode, demodel.DirectionNone, nil
		}
	case key.CodeC:
		if e.Direction == key.DirPress && isCopyModifier(e) {
			actions.CopyClipboard(position.DotStart, position.DotEnd, buff)
			return NormalMode, demodel.DirectionNone, nil
		}
	}
	return NormalMode, demodel.DirectionNone, ErrInvalid
}
<commit_msg>Add vi-like bindings for I and A.<commit_after>package kbmap

import (
	//"fmt"

	"github.com/driusan/de/actions"
	"github.com/driusan/de/demodel"
	"github.com/driusan/de/position"
	"golang.org/x/mobile/event/key"
)

var Repeat uint

func normalMap(e key.Event, buff *demodel.CharBuffer, v demodel.Viewport) (demodel.Map, demodel.ScrollDirection, error) {
	// things only happen on key press in normal mode, if it's a release
	// or a repeat, ignore it. It's not an error
	if e.Direction == key.DirRelease {
		return NormalMode, demodel.DirectionNone, nil
	}
	if buff == nil {
		return NormalMode, demodel.DirectionNone, ErrInvalid
	}
	switch e.Code {
	case key.CodeEscape:
		actions.Do("SaveOrExit", buff, v)
		return NormalMode, demodel.DirectionUp, nil
	case key.CodeDeleteBackspace:
		if e.Direction == key.DirPress {
			actions.DeleteCursor(position.DotStart, position.DotEnd, buff)
		}
		// There's a potential that we're pressing backspace at the start of the
		// screen and may need to scroll up
		return NormalMode, demodel.DirectionUp, nil
	case key.CodeI:
		if e.Modifiers&key.ModShift == key.ModShift {
			actions.MoveCursor(position.StartOfLine, position.StartOfLine, buff)
		}

		return InsertMode, demodel.DirectionNone, nil
	case key.CodeA:
		if buff.Dot.End >= uint(len(buff.Buffer)) {
			return InsertMode, demodel.DirectionDown, nil
		}

		if e.Modifiers&key.ModShift == key.ModShift {
			actions.MoveCursor(position.EndOfLine, position.EndOfLine, buff)
		} else {
			if buff.Buffer[buff.Dot.End] == '\n' {
				// vi doesn't actually let you navigate your cursor *on top*
				// of the newline, while de does.
				// It feels weird to push 'a' there and end up on the next
				// line, so if we append to the end of a line, just leave
				// the cursor end where it is,
				actions.MoveCursor(position.DotEnd, position.DotEnd, buff)
			} else {
				actions.MoveCursor(position.NextChar, position.NextChar, buff)

			}
		}

		// There's a potentially we pressed 'a' at the very end of the screen and
		// need to scroll down
		return InsertMode, demodel.DirectionDown, nil
	case key.CodeK:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				// ctrl is pressed, so just move the start and expand dot
				actions.MoveCursor(position.PrevLine, position.DotEnd, buff)
			} else {
				// ctrl is not pressed, so move the cursor without selecting
				actions.MoveCursor(position.PrevLine, position.PrevLine, buff)

			}
		}

		return NormalMode, demodel.DirectionUp, nil
	case key.CodeH:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				// ctrl is pressed, so just move the end and expand dot
				actions.MoveCursor(position.PrevChar, position.DotEnd, buff)
			} else {
				// ctrl is not pressed, so move the cursor without selecting
				actions.MoveCursor(position.PrevChar, position.PrevChar, buff)

			}
		}

		return NormalMode, demodel.DirectionUp, nil

	case key.CodeL:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				// ctrl is pressed, so just move the end and expand dot
				actions.MoveCursor(position.DotStart, position.NextChar, buff)
			} else {
				// ctrl is not pressed, so move the cursor without selecting
				actions.MoveCursor(position.NextChar, position.NextChar, buff)
			}
		}

		return NormalMode, demodel.DirectionDown, nil
	case key.CodeJ:
		if e.Modifiers&key.ModShift != 0 {
			// 'J' for Join

			undoDot := -1
			if buff.Dot.Start == buff.Dot.End {
				if Repeat <= 1 {
					Repeat = 2
				}
				undoDot = int(buff.Dot.Start)
				// Join n lines
				actions.MoveCursor(position.StartOfLine, position.DotEnd, buff)
				for ; Repeat > 1; Repeat-- {
					actions.MoveCursor(position.DotStart, position.NextLine, buff)
				}
				actions.MoveCursor(position.DotStart, position.EndOfLine, buff)
			}

			buff.JoinLines(buff.Dot.Start, buff.Dot.End)
			if undoDot >= 0 {
				buff.Undo.Dot.Start = uint(undoDot)
				buff.Undo.Dot.End = uint(undoDot)
				buff.Dot = buff.Undo.Dot
			}
			v.InvalidateCache()
			Repeat = 0
			return NormalMode, demodel.DirectionNone, nil
		}

		// 'j' for movement
		if Repeat == 0 {
			Repeat = 1

		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				actions.MoveCursor(position.DotStart, position.NextLine, buff)
			} else {
				actions.MoveCursor(position.NextLine, position.NextLine, buff)
			}
		}

		return NormalMode, demodel.DirectionDown, nil
	case key.CodeX:
		if e.Direction == key.DirPress && isCopyModifier(e) {
			actions.CutClipboard(position.DotStart, position.DotEnd, buff)
			return NormalMode, demodel.DirectionNone, nil
		}

		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			actions.MoveCursor(position.DotStart, position.NextChar, buff)
		}
		actions.DeleteCursor(position.DotStart, position.DotEnd, buff)
		return NormalMode, demodel.DirectionNone, nil
	case key.CodeW:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				actions.MoveCursor(position.DotStart, position.NextWordStart, buff)
			} else {
				actions.MoveCursor(position.NextWordStart, position.NextWordStart, buff)
			}
		}
		return NormalMode, demodel.DirectionDown, nil
	case key.CodeB:
		if Repeat == 0 {
			Repeat = 1
		}
		for ; Repeat > 0; Repeat-- {
			if e.Modifiers&key.ModControl != 0 {
				actions.MoveCursor(position.PrevWordStart, position.DotEnd, buff)
			} else {
				actions.MoveCursor(position.PrevWordStart, position.PrevWordStart, buff)
			}
		}
		return NormalMode, demodel.DirectionUp, nil
	case key.CodeP:
		actions.InsertSnarfBuffer(position.DotStart, position.DotEnd, buff)
		return NormalMode, demodel.DirectionDown, nil
	case key.CodeRightArrow:
		// Arrow keys indicate their scroll direction via the error return value,
		// they return demodel.DirectionNone to make sure both code paths don't accidentally
		// get triggered
		return NormalMode, demodel.DirectionNone, ErrScrollRight
	case key.CodeLeftArrow:
		return NormalMode, demodel.DirectionNone, ErrScrollLeft
	case key.CodeDownArrow:
		return NormalMode, demodel.DirectionNone, ErrScrollDown
	case key.CodeUpArrow:
		return NormalMode, demodel.DirectionNone, ErrScrollUp
	case key.Code0:
		if e.Modifiers == 0 {
			Repeat *= 10
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code1:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 1
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code2:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 2

		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code3:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 3

		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code4:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 4

		} else if e.Modifiers&key.ModShift != 0 {
			// $ is pressed, in most key maps..
			// it doesn't mean anything to repeat "End of line", so just
			// reset the repeat counter.
			Repeat = 0

			if e.Modifiers&key.ModControl != 0 {
				// ctrl is pressed, so just move the end and expand dot
				// BUG: This doesn't seem to work. Probably a bug in the shiny driver.
				actions.MoveCursor(position.DotStart, position.EndOfLine, buff)
			} else {
				actions.MoveCursor(position.EndOfLine, position.EndOfLine, buff)
			}

		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code5:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 5
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code6:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 6
		} else if e.Modifiers&key.ModShift != 0 {
			// ^ is pressed, on most keyboards..
			// It doesn't mean anything to repeat "Start of line" either
			Repeat = 0

			if e.Modifiers&key.ModControl != 0 {
				// BUG: This doesn't seem to work. Probably a bug in the shiny driver.
				actions.MoveCursor(position.StartOfLine, position.DotEnd, buff)
			} else {
				actions.MoveCursor(position.StartOfLine, position.StartOfLine, buff)
			}

		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code7:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 7
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code8:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 8
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.Code9:
		if e.Modifiers == 0 {
			Repeat = Repeat*10 + 9
		}
		return NormalMode, demodel.DirectionNone, nil
	case key.CodeG:
		if e.Modifiers&key.ModShift != 0 {
			// G treats the repeat counter differently. If it's set, it means "Go to
			// lineno", and if it's not it means "Go to end of buffer"
			if Repeat > 0 {
				/* This seems to be buggy, so do it manually, at least until there's better
				   tests.
				actions.MoveCursor(position.BuffStart, position.BuffStart, buff)
				for ; Repeat > 0; Repeat-- {
					actions.MoveCursor(position.NextLine, position.NextLine, buff)
				}*/
				lineNo := uint(1)
				for i, c := range buff.Buffer {
					if c == '\n' {
						lineNo++
					}
					// should be equal, but if we've went too far at least it's
					// better to stop now.
					if lineNo >= Repeat {
						// we're on the \n itself, so add 1 to get to the actual line.
						if Repeat == 1 {
							buff.Dot.Start = 0
						} else {
							buff.Dot.Start = uint(i + 1)
						}
						buff.Dot.End = buff.Dot.Start
						Repeat = 0
						// select the line.
						actions.MoveCursor(position.DotStart, position.EndOfLine, buff)
						break
					}
				}
				// We don't know if the line being moved to is above or below, but send an "up"
				// hint so that it centers it at the top of the screen based on dot.
				return NormalMode, demodel.DirectionUp, nil
			}
			// moving to the end of the file because Repeat isn't set.
			actions.MoveCursor(position.BuffEnd, position.BuffEnd, buff)
			// send a down hint so it centers it at the bottom.
			return NormalMode, demodel.DirectionDown, nil
		}

	case key.CodeD:
		return DeleteMode, demodel.DirectionNone, nil
	case key.CodeU:
		actions.Do("Undo", buff, v)
		return NormalMode, demodel.DirectionNone, nil
	case key.CodeO:
		actions.MoveCursor(position.EndOfLineWithNewline, position.EndOfLineWithNewline, buff)

		newBuffer := make([]byte, len(buff.Buffer)+1)
		copy(newBuffer, buff.Buffer[:buff.Dot.Start])
		newBuffer[buff.Dot.Start] = '\n'
		copy(newBuffer[buff.Dot.Start+1:], buff.Buffer[buff.Dot.End:])

		buff.Buffer = newBuffer
		return InsertMode, demodel.DirectionDown, nil
	case key.CodeSlash:
		if buff.Dot.Start == buff.Dot.End {
			actions.FindNext(position.CurWordStart, position.CurWordEnd, buff)

		} else {
			actions.FindNext(position.DotStart, position.DotEnd, buff)
		}
		return NormalMode, demodel.DirectionDown, nil
	case key.CodeReturnEnter:
		if buff.Dot.Start == buff.Dot.End {
			actions.OpenOrPerformAction(position.CurExecutionWordStart, position.CurExecutionWordEnd, buff, v)
		} else {
			actions.OpenOrPerformAction(position.DotStart, position.DotEnd, buff, v)
		}
		// There's a possibility OpenOrPerformAction opened a new file, in which case
		// we should scroll to the top, or inserted text, in which case we should scroll
		// to the top of the inserted text. Either way, it's an "Up" hint so that the
		// scrolling is based on Dot.Start
		return v.GetKeyboardMode(), demodel.DirectionUp, nil
	case key.CodeSemicolon:
		if e.Modifiers&key.ModShift != 0 {
			if buff.Tagline != nil {
				buff.Tagline.Buffer = append(buff.Tagline.Buffer, ' ')

				buff.Tagline.Dot.Start = uint(len(buff.Tagline.Buffer))
				buff.Tagline.Dot.End = buff.Tagline.Dot.Start
			}
		}
		return TagMode, demodel.DirectionNone, nil
	case key.CodeV:
		if e.Direction == key.DirPress && isCopyModifier(e) {
			actions.PasteClipboard(position.DotStart, position.DotEnd, buff)
			return NormalMode, demodel.DirectionNone, nil
		}
	case key.CodeC:
		if e.Direction == key.DirPress && isCopyModifier(e) {
			actions.CopyClipboard(position.DotStart, position.DotEnd, buff)
			return NormalMode, demodel.DirectionNone, nil
		}
	}
	return NormalMode, demodel.DirectionNone, ErrInvalid
}
<|endoftext|>


<commit_before>// Copyright 2020 The LUCI Authors.
//
// 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 realmset provides queryable representation of LUCI Realms DB.
//
// Used internally by authdb.Snapshot.
package realmset

import (
	"strings"

	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"

	"go.chromium.org/luci/server/auth/realms"
	"go.chromium.org/luci/server/auth/service/protocol"

	"go.chromium.org/luci/server/auth/authdb/internal/graph"
)

// ExpectedAPIVersion is the supported value of api_version field.
//
// See Build implementation for details.
const ExpectedAPIVersion = 1

// Realms is a queryable representation of realms.Realms proto.
type Realms struct {
	perms  map[string]PermissionIndex       // permission name -> its index
	names  stringset.Set                    // just names of all defined realms
	realms map[realmAndPerm]groupsAndIdents // <realm, perm> -> who has it
	data   map[string]*protocol.RealmData   // per-realm attached RealmData
}

// PermissionIndex is used in place of permission names.
//
// Note: should match an int type used in `permissions` field in the proto.
type PermissionIndex uint32

// Condition is a predicate on attributes map.
//
// `nil` condition evaluates to true.
type Condition struct{}

// Eval evaluates the condition over given attributes.
func (c *Condition) Eval(attrs realms.Attrs) bool {
	// TODO(vadimsh): Implement.
	return true
}

// Binding represents a set of principals and a condition when it can be used.
//
// See Bindings(...) method for more details.
type Binding struct {
	Condition *Condition
	Groups    graph.SortedNodeSet
	Idents    stringset.Set
}

// realmAndPerm is used as a composite key in `realms` map.
type realmAndPerm struct {
	realm string
	perm  PermissionIndex
}

// groupsAndIdents is used as a value in `realms` map.
type groupsAndIdents struct {
	groups graph.SortedNodeSet
	idents stringset.Set
}

// PermissionIndex returns an index of the given permission.
//
// It can be passed to Bindings(...). Returns (0, false) if there's no such
// permission in the Realms DB.
func (r *Realms) PermissionIndex(perm realms.Permission) (idx PermissionIndex, ok bool) {
	idx, ok = r.perms[perm.Name()]
	return
}

// HasRealm returns true if the given realm exists in the DB.
func (r *Realms) HasRealm(realm string) bool {
	return r.names.Has(realm)
}

// Data returns RealmData attached to a realm or nil if none.
func (r *Realms) Data(realm string) *protocol.RealmData {
	return r.data[realm]
}

// Bindings returns representation of bindings that define who has the requested
// permission in the given realm.
//
// Each returned binding is a tuple (condition, groups, identities):
//   * Condition: a predicate over realms.Attrs map that evaluates to true if
//     this binding is "active". Inactive bindings should be skipped.
//   * Groups: a set of groups with principals that have the permission,
//     represented by a sorted slice of group indexes in a graph.QueryableGraph
//     which was passed to Build().
//   * Identities: a set of identity strings that were specified in the realm
//     ACL directly (not via a group).
//
// The permission should be specified as its index obtained via PermissionIndex.
//
// The realm name is not validated. Unknown or invalid realms are silently
// treated as empty. No fallback to @root happens.
//
// Returns nil if the requested permission is not mentioned in any binding in
// the realm at all.
func (r *Realms) Bindings(realm string, perm PermissionIndex) []Binding {
	if out, ok := r.realms[realmAndPerm{realm, perm}]; ok {
		return []Binding{{nil, out.groups, out.idents}}
	} else {
		return nil
	}
}

// Build constructs Realms from the proto message and the group graph.
func Build(r *protocol.Realms, qg *graph.QueryableGraph) (*Realms, error) {
	// Do not use realms.Realms we don't understand. Better to go offline
	// completely than mistakenly allow access to something private by
	// misinterpreting realm rules (e.g. if a new hypothetical DENY rule is
	// misinterpreted as ALLOW).
	//
	// Bumping `api_version` (if it ever happens) should be done extremely
	// carefully in multiple stages:
	//   1. Update components.auth to understand both new and old api_version.
	//   2. Redeploy *everything*.
	//   3. Update Auth Service to generate realms.Realms using the new API.
	if r.ApiVersion != ExpectedAPIVersion {
		return nil, errors.Reason(
			"Realms proto has api_version %d not compatible with this service (it expects %d)",
			r.ApiVersion, ExpectedAPIVersion).Err()
	}

	// Build map: permission name -> its index (since Binding messages operate
	// with indexes). Using ints as keys is also slightly faster than strings.
	perms := make(map[string]PermissionIndex, len(r.Permissions))
	for idx, perm := range r.Permissions {
		perms[perm.Name] = PermissionIndex(idx)
	}

	// Gather names of all realms for HasRealm check.
	names := stringset.New(len(r.Realms))
	for _, realm := range r.Realms {
		names.Add(realm.Name)
	}

	// This is the `realms` map under construction. We'll shrink its memory
	// footprint at the end. Just like `realms` it uses a composite key
	// realmAndPerm as a more memory-efficient alternative to a map of maps
	// (realm -> perm -> principals).
	realmsToBe := map[realmAndPerm]principalSet{}

	// interner is used to deduplicate memory used to store identity names.
	interner := stringInterner{}

	// Visit all bindings in all realms and update principal sets in realmsToBe.
	for _, realm := range r.Realms {
		for _, binding := range realm.Bindings {
			// Categorize 'principals' into groups and identity strings.
			groups, idents := categorizePrincipals(binding.Principals, qg, interner)
			if len(groups) == 0 && len(idents) == 0 {
				continue
			}

			// Add them into the corresponding principal sets in realmsToBe.
			for _, permIdx := range binding.Permissions {
				key := realmAndPerm{realm.Name, PermissionIndex(permIdx)}
				if ps, ok := realmsToBe[key]; ok {
					ps.add(groups, idents)
				} else {
					realmsToBe[key] = newPrincipalSet(groups, idents)
				}
			}
		}
	}

	// Replace identically looking group sets with references to a single copy.
	realms := make(map[realmAndPerm]groupsAndIdents, len(realmsToBe))
	dedupper := graph.NodeSetDedupper{}
	for key, ps := range realmsToBe {
		realms[key] = ps.finalize(dedupper)
	}

	// Extract attached per-realm data into a queryable map.
	count := 0
	for _, realm := range r.Realms {
		if realm.Data != nil {
			count++
		}
	}
	data := make(map[string]*protocol.RealmData, count)
	for _, realm := range r.Realms {
		if realm.Data != nil {
			data[realm.Name] = realm.Data
		}
	}

	return &Realms{
		perms:  perms,
		names:  names,
		realms: realms,
		data:   data,
	}, nil
}

// stringInterner implements string interning to save some memory.
type stringInterner map[string]string

// intern returns an interned copy of 's'.
func (si stringInterner) intern(s string) string {
	if existing, ok := si[s]; ok {
		return existing
	}
	si[s] = s
	return s
}

// categorizePrincipals splits a list of principals into a list of groups
// (identified by their indexes in a QueryableGraph) and list of identity names.
//
// Unknown groups are silently skipped.
func categorizePrincipals(p []string, qg *graph.QueryableGraph, interner stringInterner) (groups []graph.NodeIndex, idents []string) {
	for _, principal := range p {
		if strings.HasPrefix(principal, "group:") {
			if idx, ok := qg.GroupIndex(strings.TrimPrefix(principal, "group:")); ok {
				groups = append(groups, idx)
			}
		} else {
			idents = append(idents, interner.intern(principal))
		}
	}
	return
}

// principalSet represents a set of groups and identities.
//
// It is used transiently when constructing the final memory-optimized set in
// Build.
type principalSet struct {
	groups graph.NodeSet
	idents stringset.Set
}

func newPrincipalSet(groups []graph.NodeIndex, idents []string) principalSet {
	ps := principalSet{
		groups: make(graph.NodeSet, len(groups)),
		idents: stringset.New(len(idents)),
	}
	ps.add(groups, idents)
	return ps
}

func (ps principalSet) add(groups []graph.NodeIndex, idents []string) {
	for _, idx := range groups {
		ps.groups.Add(idx)
	}
	for _, ident := range idents {
		ps.idents.Add(ident)
	}
}

// finalize produces a memory-optimized representation of this principal set.
//
// It replace identically looking group sets with references to a single copy
// using the given dedupper. It also throws away zero-length sets replacing them
// with nils.
//
// Non-empty identity sets are kept as is without any dedupping, assuming using
// identities in Realm ACLs directly is rare and not worth optimizing for (on
// top of the string interning optimization we've already done).
func (ps principalSet) finalize(dedupper graph.NodeSetDedupper) groupsAndIdents {
	var groups graph.SortedNodeSet
	if len(ps.groups) > 0 {
		groups = dedupper.Dedup(ps.groups)
	}
	var idents stringset.Set
	if ps.idents.Len() > 0 {
		idents = ps.idents
	}
	return groupsAndIdents{groups, idents}
}
<commit_msg>[realms] Make realmset.Build aware of conditions.<commit_after>// Copyright 2020 The LUCI Authors.
//
// 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 realmset provides queryable representation of LUCI Realms DB.
//
// Used internally by authdb.Snapshot.
package realmset

import (
	"sort"
	"strings"

	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"

	"go.chromium.org/luci/server/auth/realms"
	"go.chromium.org/luci/server/auth/service/protocol"

	"go.chromium.org/luci/server/auth/authdb/internal/graph"
)

// ExpectedAPIVersion is the supported value of api_version field.
//
// See Build implementation for details.
const ExpectedAPIVersion = 1

// Realms is a queryable representation of realms.Realms proto.
type Realms struct {
	perms  map[string]PermissionIndex     // permission name -> its index
	names  stringset.Set                  // just names of all defined realms
	realms map[realmAndPerm][]Binding     // <realm, perm> -> who has it under which conditions
	data   map[string]*protocol.RealmData // per-realm attached RealmData
}

// PermissionIndex is used in place of permission names.
//
// Note: should match an int type used in `permissions` field in the proto.
type PermissionIndex uint32

// Condition is a predicate on attributes map.
//
// `nil` condition evaluates to true.
type Condition struct{}

// Eval evaluates the condition over given attributes.
func (c *Condition) Eval(attrs realms.Attrs) bool {
	// TODO(vadimsh): Implement.
	return true
}

// Binding represents a set of principals and a condition when it can be used.
//
// See Bindings(...) method for more details.
type Binding struct {
	Condition *Condition
	Groups    graph.SortedNodeSet
	Idents    stringset.Set
}

// badness is an overall heuristic score of how complex this binding to
// evaluate and how likely it will apply.
//
// 0 means "easy to evaluate, high likelihood of applying". Used to sort
// bindings in []Binding array returned by Bindings(...).
func (b *Binding) badness() int {
	// TODO(vadimsh): This can be improved. For example, bindings with groups
	// that contain globs (like `user:*`) are more likely to apply and should
	// have lower badness.
	if b.Condition == nil {
		return 0
	}
	return 1
}

// realmAndPerm is used as a composite key in `realms` map.
type realmAndPerm struct {
	realm string
	perm  PermissionIndex
}

// PermissionIndex returns an index of the given permission.
//
// It can be passed to Bindings(...). Returns (0, false) if there's no such
// permission in the Realms DB.
func (r *Realms) PermissionIndex(perm realms.Permission) (idx PermissionIndex, ok bool) {
	idx, ok = r.perms[perm.Name()]
	return
}

// HasRealm returns true if the given realm exists in the DB.
func (r *Realms) HasRealm(realm string) bool {
	return r.names.Has(realm)
}

// Data returns RealmData attached to a realm or nil if none.
func (r *Realms) Data(realm string) *protocol.RealmData {
	return r.data[realm]
}

// Bindings returns representation of bindings that define who has the requested
// permission in the given realm.
//
// Each returned binding is a tuple (condition, groups, identities):
//   * Condition: a predicate over realms.Attrs map that evaluates to true if
//     this binding is "active". Inactive bindings should be skipped.
//   * Groups: a set of groups with principals that have the permission,
//     represented by a sorted slice of group indexes in a graph.QueryableGraph
//     which was passed to Build().
//   * Identities: a set of identity strings that were specified in the realm
//     ACL directly (not via a group).
//
// The permission should be specified as its index obtained via PermissionIndex.
//
// The realm name is not validated. Unknown or invalid realms are silently
// treated as empty. No fallback to @root happens.
//
// Returns nil if the requested permission is not mentioned in any binding in
// the realm at all.
func (r *Realms) Bindings(realm string, perm PermissionIndex) []Binding {
	return r.realms[realmAndPerm{realm, perm}]
}

// Build constructs Realms from the proto message and the group graph.
func Build(r *protocol.Realms, qg *graph.QueryableGraph) (*Realms, error) {
	// Do not use realms.Realms we don't understand. Better to go offline
	// completely than mistakenly allow access to something private by
	// misinterpreting realm rules (e.g. if a new hypothetical DENY rule is
	// misinterpreted as ALLOW).
	//
	// Bumping `api_version` (if it ever happens) should be done extremely
	// carefully in multiple stages:
	//   1. Update components.auth to understand both new and old api_version.
	//   2. Redeploy *everything*.
	//   3. Update Auth Service to generate realms.Realms using the new API.
	if r.ApiVersion != ExpectedAPIVersion {
		return nil, errors.Reason(
			"Realms proto has api_version %d not compatible with this service (it expects %d)",
			r.ApiVersion, ExpectedAPIVersion).Err()
	}

	// Build map: permission name -> its index (since Binding messages operate
	// with indexes). Using ints as keys is also slightly faster than strings.
	perms := make(map[string]PermissionIndex, len(r.Permissions))
	for idx, perm := range r.Permissions {
		perms[perm.Name] = PermissionIndex(idx)
	}

	// Gather names of all realms for HasRealm check.
	names := stringset.New(len(r.Realms))
	for _, realm := range r.Realms {
		names.Add(realm.Name)
	}

	// This is the `realms` map under construction. We'll shrink its memory
	// footprint at the end.
	type bindingKey struct {
		realmAndPerm
		cond *Condition
	}
	realmsToBe := map[bindingKey]principalSet{}
	counts := map[realmAndPerm]int{}

	// interner is used to deduplicate memory used to store identity names.
	interner := stringInterner{}

	// Visit all bindings in all realms and update principal sets in realmsToBe.
	for _, realm := range r.Realms {
		for _, binding := range realm.Bindings {
			// Categorize 'principals' into groups and identity strings.
			groups, idents := categorizePrincipals(binding.Principals, qg, interner)
			if len(groups) == 0 && len(idents) == 0 {
				continue
			}

			// Build a condition predicate. `nil` means "no condition".
			//
			// TODO.
			var cond *Condition

			// Add them into the corresponding principal sets in realmsToBe.
			for _, permIdx := range binding.Permissions {
				key := bindingKey{realmAndPerm{realm.Name, PermissionIndex(permIdx)}, cond}
				if ps, ok := realmsToBe[key]; ok {
					ps.add(groups, idents)
				} else {
					realmsToBe[key] = newPrincipalSet(groups, idents)
					counts[key.realmAndPerm] += 1
				}
			}
		}
	}

	// Replace identically looking group sets with references to a single copy.
	// Collect conditional bindings for the same (realm, perm) key into an array,
	// since we'll need to evaluate them sequentially when serving HasPermission
	// checks.
	realms := make(map[realmAndPerm][]Binding, len(counts))
	dedupper := graph.NodeSetDedupper{}
	for key, ps := range realmsToBe {
		groups, idents := ps.finalize(dedupper)
		if realms[key.realmAndPerm] == nil {
			realms[key.realmAndPerm] = make([]Binding, 0, counts[key.realmAndPerm])
		}
		realms[key.realmAndPerm] = append(realms[key.realmAndPerm], Binding{
			Condition: key.cond,
			Groups:    groups,
			Idents:    idents,
		})
	}

	// Order bindings by "badness" of checking (easiest to check first) and
	// chances of applying. Right now this uses a very simplistic heuristic:
	// unconditional bindings are easier to check and more likely to apply than
	// conditional ones.
	for _, bindings := range realms {
		sort.Slice(bindings, func(i, j int) bool {
			return bindings[i].badness() < bindings[j].badness()
		})
	}

	// Extract attached per-realm data into a queryable map.
	count := 0
	for _, realm := range r.Realms {
		if realm.Data != nil {
			count++
		}
	}
	data := make(map[string]*protocol.RealmData, count)
	for _, realm := range r.Realms {
		if realm.Data != nil {
			data[realm.Name] = realm.Data
		}
	}

	return &Realms{
		perms:  perms,
		names:  names,
		realms: realms,
		data:   data,
	}, nil
}

// stringInterner implements string interning to save some memory.
type stringInterner map[string]string

// intern returns an interned copy of 's'.
func (si stringInterner) intern(s string) string {
	if existing, ok := si[s]; ok {
		return existing
	}
	si[s] = s
	return s
}

// categorizePrincipals splits a list of principals into a list of groups
// (identified by their indexes in a QueryableGraph) and list of identity names.
//
// Unknown groups are silently skipped.
func categorizePrincipals(p []string, qg *graph.QueryableGraph, interner stringInterner) (groups []graph.NodeIndex, idents []string) {
	for _, principal := range p {
		if strings.HasPrefix(principal, "group:") {
			if idx, ok := qg.GroupIndex(strings.TrimPrefix(principal, "group:")); ok {
				groups = append(groups, idx)
			}
		} else {
			idents = append(idents, interner.intern(principal))
		}
	}
	return
}

// principalSet represents a set of groups and identities.
//
// It is used transiently when constructing the final memory-optimized set in
// Build.
type principalSet struct {
	groups graph.NodeSet
	idents stringset.Set
}

func newPrincipalSet(groups []graph.NodeIndex, idents []string) principalSet {
	ps := principalSet{
		groups: make(graph.NodeSet, len(groups)),
		idents: stringset.New(len(idents)),
	}
	ps.add(groups, idents)
	return ps
}

func (ps principalSet) add(groups []graph.NodeIndex, idents []string) {
	for _, idx := range groups {
		ps.groups.Add(idx)
	}
	for _, ident := range idents {
		ps.idents.Add(ident)
	}
}

// finalize produces a memory-optimized representation of this principal set.
//
// It replace identically looking group sets with references to a single copy
// using the given dedupper. It also throws away zero-length sets replacing them
// with nils.
//
// Non-empty identity sets are kept as is without any dedupping, assuming using
// identities in Realm ACLs directly is rare and not worth optimizing for (on
// top of the string interning optimization we've already done).
func (ps principalSet) finalize(dedupper graph.NodeSetDedupper) (graph.SortedNodeSet, stringset.Set) {
	var groups graph.SortedNodeSet
	if len(ps.groups) > 0 {
		groups = dedupper.Dedup(ps.groups)
	}
	var idents stringset.Set
	if ps.idents.Len() > 0 {
		idents = ps.idents
	}
	return groups, idents
}
<|endoftext|>


<commit_before>package crawl

import (
	"io"
	"net/http"
	"net/http/cookiejar"
	"strings"
	"sync"
	"time"

	"github.com/ryanuber/go-glob"

	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

// Handler - Crawler handler.
type Handler func(context.Context, *Response) error

// Middleware - Crawler middleware.
type Middleware func(context.Context, *Request, *http.Request) error

// Crawler - Crawler interface.
type Crawler interface {
	// Schedule - Schedules request.
	// Context is passed to queue in a job.
	Schedule(context.Context, *Request) error

	// Execute - Makes a http request respecting context deadline.
	// If request Raw is not true - ParseHTML() method is executed on Response.
	// Then all callbacks are executed with context.
	Execute(context.Context, *Request) (*Response, error)

	// Handlers - Returns all registered handlers.
	Handlers() map[string][]Handler

	// Register - Registers crawl handler.
	Register(name string, h Handler)

	// Middleware - Registers a middleware.
	// Request is not executed if middleware returns an error.
	Middleware(Middleware)

	// Start - Starts the crawler.
	// Crawler will work until queue is empty or closed.
	// All errors should be received from Errors() channel.
	Start()

	// Errors - Returns channel that will receive all crawl errors.
	// Only errors from queued requests are here.
	// Not only request errors but also queue errors.
	Errors() <-chan error
}

// New - Creates new crawler. If queue is not provided it will create
// a memory queue with a capacity of WithQueueCapacity seting value (default=10000).
func New(opts ...Option) Crawler {
	c := &crawl{
		handlers:   make(map[string][]Handler),
		errorsChan: make(chan error, 10000),
		opts: &options{
			concurrency:   1000,
			queueCapacity: 10000,
			headers:       DefaultHeaders,
		},
	}
	for _, opt := range opts {
		opt(c)
	}
	if c.transport == nil {
		c.transport = new(http.Transport)
	}
	c.client = &http.Client{Transport: c.transport}
	if c.client.Jar == nil {
		c.client.Jar, _ = cookiejar.New(nil)
	}
	if c.queue == nil {
		c.queue = NewQueue(c.opts.queueCapacity)
	}
	return c
}

// DefaultHeaders - Default crawler headers.
var DefaultHeaders = map[string]string{
	"Accept":          "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
	"Accept-Language": "en-US,en;q=0.8",
	"User-Agent":      "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/46.0.2490.86 Safari/537.36",
}

// crawl - Crawler structure.
type crawl struct {
	errorsChan chan error
	handlers   map[string][]Handler
	transport  *http.Transport
	client     *http.Client
	opts       *options

	queue Queue

	// patterns - callbacks glob patterns
	patterns []string

	// middlewares - crawler middlewares.
	middlewares []Middleware
}

func (crawl *crawl) Start() {
	wg := new(sync.WaitGroup)
	for i := 0; i < crawl.opts.concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				job, err := crawl.queue.Get()
				if err == io.EOF {
					return
				} else if err != nil {
					crawl.errorsChan <- err
					return
				}

				if _, err := crawl.Execute(job.Context(), job.Request()); err != nil {
					crawl.errorsChan <- &RequestError{Err: err, Request: job.Request()}
				}

				job.Done()
			}
		}()
	}
	wg.Wait()
	return
}

func (crawl *crawl) Execute(ctx context.Context, req *Request) (resp *Response, err error) {
	// Set default timeout if enabled and empty in context
	if _, ok := ctx.Deadline(); !ok && crawl.opts.defaultTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithDeadline(ctx, time.Now().Add(crawl.opts.defaultTimeout))
		defer cancel()
	}

	// Get http.Request structure
	httpReq, err := ConstructHTTPRequest(req)
	if err != nil {
		return
	}

	// Run request middlewares
	for _, middleware := range crawl.middlewares {
		if err = middleware(ctx, req, httpReq); err != nil {
			return
		}
	}

	// Copy default headers
	for name, value := range crawl.opts.headers {
		if _, has := httpReq.Header[name]; !has {
			httpReq.Header.Set(name, value)
		}
	}

	httpResp, err := ctxhttp.Do(ctx, crawl.client, httpReq)
	if err != nil {
		return
	}

	resp = &Response{
		Response: httpResp,
		Request:  req,
	}
	defer resp.Close()

	// Parse HTML if not request.Raw
	if !req.Raw {
		err = resp.ParseHTML()
		if err != nil {
			return
		}
	}

	if err = crawl.executeHandlers(ctx, resp); err != nil {
		return nil, err
	}

	return
}

func (crawl *crawl) executeHandlers(ctx context.Context, resp *Response) (err error) {
	handlers := crawl.getHandlers(resp.Request.Callbacks)
	if len(handlers) == 0 {
		return
	}
	for _, handler := range handlers {
		if err = handler(ctx, resp); err != nil {
			return
		}
	}
	return
}

func (crawl *crawl) getHandlers(callbacks []string) (list []Handler) {
	for _, pattern := range crawl.patterns {
		for _, name := range callbacks {
			if glob.Glob(pattern, name) {
				list = append(list, crawl.handlers[pattern]...)
				break
			}
		}
	}
	for _, name := range callbacks {
		list = append(list, crawl.handlers[name]...)
	}
	return
}

func (crawl *crawl) Middleware(m Middleware) {
	crawl.middlewares = append(crawl.middlewares, m)
}

func (crawl *crawl) Register(name string, h Handler) {
	if _, ok := crawl.handlers[name]; !ok && strings.Contains(name, "*") {
		crawl.patterns = append(crawl.patterns, name)
	}
	crawl.handlers[name] = append(crawl.handlers[name], h)
}

func (crawl *crawl) Schedule(ctx context.Context, req *Request) error {
	return crawl.queue.Schedule(ctx, req)
}

func (crawl *crawl) Errors() <-chan error {
	return crawl.errorsChan
}

func (crawl *crawl) Handlers() map[string][]Handler {
	return crawl.handlers
}
<commit_msg>proxy: use proxy from context<commit_after>package crawl

import (
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/ryanuber/go-glob"

	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

// Handler - Crawler handler.
type Handler func(context.Context, *Response) error

// Middleware - Crawler middleware.
type Middleware func(context.Context, *Request, *http.Request) error

// Crawler - Crawler interface.
type Crawler interface {
	// Schedule - Schedules request.
	// Context is passed to queue in a job.
	Schedule(context.Context, *Request) error

	// Execute - Makes a http request respecting context deadline.
	// If request Raw is not true - ParseHTML() method is executed on Response.
	// Then all callbacks are executed with context.
	Execute(context.Context, *Request) (*Response, error)

	// Handlers - Returns all registered handlers.
	Handlers() map[string][]Handler

	// Register - Registers crawl handler.
	Register(name string, h Handler)

	// Middleware - Registers a middleware.
	// Request is not executed if middleware returns an error.
	Middleware(Middleware)

	// Start - Starts the crawler.
	// Crawler will work until queue is empty or closed.
	// All errors should be received from Errors() channel.
	Start()

	// Errors - Returns channel that will receive all crawl errors.
	// Only errors from queued requests are here.
	// Not only request errors but also queue errors.
	Errors() <-chan error
}

// New - Creates new crawler. If queue is not provided it will create
// a memory queue with a capacity of WithQueueCapacity seting value (default=10000).
func New(opts ...Option) Crawler {
	c := &crawl{
		handlers:   make(map[string][]Handler),
		errorsChan: make(chan error, 10000),
		opts: &options{
			concurrency:   1000,
			queueCapacity: 10000,
			headers:       DefaultHeaders,
		},
	}
	for _, opt := range opts {
		opt(c)
	}
	if c.transport == nil {
		c.transport = new(http.Transport)
	}
	c.client = &http.Client{Transport: c.transport}
	if c.client.Jar == nil {
		c.client.Jar, _ = cookiejar.New(nil)
	}
	if c.queue == nil {
		c.queue = NewQueue(c.opts.queueCapacity)
	}
	return c
}

// DefaultHeaders - Default crawler headers.
var DefaultHeaders = map[string]string{
	"Accept":          "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
	"Accept-Language": "en-US,en;q=0.8",
	"User-Agent":      "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/46.0.2490.86 Safari/537.36",
}

// crawl - Crawler structure.
type crawl struct {
	errorsChan chan error
	handlers   map[string][]Handler
	transport  *http.Transport
	client     *http.Client
	opts       *options

	queue Queue

	// patterns - callbacks glob patterns
	patterns []string

	// middlewares - crawler middlewares.
	middlewares []Middleware
}

func (crawl *crawl) Start() {
	wg := new(sync.WaitGroup)
	for i := 0; i < crawl.opts.concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				job, err := crawl.queue.Get()
				if err == io.EOF {
					return
				} else if err != nil {
					crawl.errorsChan <- err
					return
				}

				if _, err := crawl.Execute(job.Context(), job.Request()); err != nil {
					crawl.errorsChan <- &RequestError{Err: err, Request: job.Request()}
				}

				job.Done()
			}
		}()
	}
	wg.Wait()
	return
}

func (crawl *crawl) Execute(ctx context.Context, req *Request) (resp *Response, err error) {
	// Set default timeout if enabled and empty in context
	if _, ok := ctx.Deadline(); !ok && crawl.opts.defaultTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithDeadline(ctx, time.Now().Add(crawl.opts.defaultTimeout))
		defer cancel()
	}

	// Get http.Request structure
	httpReq, err := ConstructHTTPRequest(req)
	if err != nil {
		return
	}

	// Run request middlewares
	for _, middleware := range crawl.middlewares {
		if err = middleware(ctx, req, httpReq); err != nil {
			return
		}
	}

	// Copy default headers
	for name, value := range crawl.opts.headers {
		if _, has := httpReq.Header[name]; !has {
			httpReq.Header.Set(name, value)
		}
	}

	client := crawl.client
	if addrs, ok := ProxyFromContext(ctx); ok && len(addrs) > 0 {
		t := &http.Transport{
			Proxy: func(_ *http.Request) (*url.URL, error) {
				addr := addrs[rand.Intn(len(addrs))]
				u, err := url.Parse(addr)
				if err != nil || !strings.HasPrefix(u.Scheme, "http") {
					u, err = url.Parse("http://" + addr)
				}
				if err != nil {
					return nil, fmt.Errorf("invalid proxy address %q: %v", addr, err)
				}
				return u, nil
			},
		}
		client = &http.Client{Transport: t}
	}

	httpResp, err := ctxhttp.Do(ctx, client, httpReq)
	if err != nil {
		return
	}

	resp = &Response{
		Response: httpResp,
		Request:  req,
	}
	defer resp.Close()

	// Parse HTML if not request.Raw
	if !req.Raw {
		err = resp.ParseHTML()
		if err != nil {
			return
		}
	}

	if err = crawl.executeHandlers(ctx, resp); err != nil {
		return nil, err
	}

	return
}

func (crawl *crawl) executeHandlers(ctx context.Context, resp *Response) (err error) {
	handlers := crawl.getHandlers(resp.Request.Callbacks)
	if len(handlers) == 0 {
		return
	}
	for _, handler := range handlers {
		if err = handler(ctx, resp); err != nil {
			return
		}
	}
	return
}

func (crawl *crawl) getHandlers(callbacks []string) (list []Handler) {
	for _, pattern := range crawl.patterns {
		for _, name := range callbacks {
			if glob.Glob(pattern, name) {
				list = append(list, crawl.handlers[pattern]...)
				break
			}
		}
	}
	for _, name := range callbacks {
		list = append(list, crawl.handlers[name]...)
	}
	return
}

func (crawl *crawl) Middleware(m Middleware) {
	crawl.middlewares = append(crawl.middlewares, m)
}

func (crawl *crawl) Register(name string, h Handler) {
	if _, ok := crawl.handlers[name]; !ok && strings.Contains(name, "*") {
		crawl.patterns = append(crawl.patterns, name)
	}
	crawl.handlers[name] = append(crawl.handlers[name], h)
}

func (crawl *crawl) Schedule(ctx context.Context, req *Request) error {
	return crawl.queue.Schedule(ctx, req)
}

func (crawl *crawl) Errors() <-chan error {
	return crawl.errorsChan
}

func (crawl *crawl) Handlers() map[string][]Handler {
	return crawl.handlers
}
<|endoftext|>