1/*
2   Copyright The containerd Authors.
3
4   Licensed under the Apache License, Version 2.0 (the "License");
5   you may not use this file except in compliance with the License.
6   You may obtain a copy of the License at
7
8       http://www.apache.org/licenses/LICENSE-2.0
9
10   Unless required by applicable law or agreed to in writing, software
11   distributed under the License is distributed on an "AS IS" BASIS,
12   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13   See the License for the specific language governing permissions and
14   limitations under the License.
15*/
16
17package progress
18
19import (
20	"bytes"
21	"fmt"
22	"io"
23	"os"
24	"regexp"
25	"strings"
26
27	"github.com/containerd/console"
28)
29
30var (
31	regexCleanLine = regexp.MustCompile("\x1b\\[[0-9]+m[\x1b]?")
32)
33
34// Writer buffers writes until flush, at which time the last screen is cleared
35// and the current buffer contents are written. This is useful for
36// implementing progress displays, such as those implemented in docker and
37// git.
38type Writer struct {
39	buf   bytes.Buffer
40	w     io.Writer
41	lines int
42}
43
44// NewWriter returns a writer
45func NewWriter(w io.Writer) *Writer {
46	return &Writer{
47		w: w,
48	}
49}
50
51// Write the provided bytes
52func (w *Writer) Write(p []byte) (n int, err error) {
53	return w.buf.Write(p)
54}
55
56// Flush should be called when refreshing the current display.
57func (w *Writer) Flush() error {
58	if w.buf.Len() == 0 {
59		return nil
60	}
61
62	if err := w.clearLines(); err != nil {
63		return err
64	}
65	w.lines = countLines(w.buf.String())
66
67	if _, err := w.w.Write(w.buf.Bytes()); err != nil {
68		return err
69	}
70
71	w.buf.Reset()
72	return nil
73}
74
75// TODO(stevvooe): The following are system specific. Break these out if we
76// decide to build this package further.
77
78func (w *Writer) clearLines() error {
79	for i := 0; i < w.lines; i++ {
80		if _, err := fmt.Fprintf(w.w, "\x1b[1A\x1b[2K\r"); err != nil {
81			return err
82		}
83	}
84
85	return nil
86}
87
88// countLines in the output. If a line is longer than the console width then
89// an extra line is added to the count for each wrapped line. If the console
90// width is undefined then 0 is returned so that no lines are cleared on the next
91// flush.
92func countLines(output string) int {
93	con, err := console.ConsoleFromFile(os.Stdin)
94	if err != nil {
95		return 0
96	}
97	ws, err := con.Size()
98	if err != nil {
99		return 0
100	}
101	width := int(ws.Width)
102	if width <= 0 {
103		return 0
104	}
105	strlines := strings.Split(output, "\n")
106	lines := -1
107	for _, line := range strlines {
108		lines += (len(stripLine(line))-1)/width + 1
109	}
110	return lines
111}
112
113func stripLine(line string) string {
114	return string(regexCleanLine.ReplaceAll([]byte(line), []byte{}))
115}
116