github.com/orangebees/go-oneutils@v0.0.10/Compress/func.go (about)

     1  package Compress
     2  
     3  import (
     4  	"archive/tar"
     5  	"bytes"
     6  	"github.com/valyala/bytebufferpool"
     7  	"github.com/valyala/fasthttp"
     8  	"io"
     9  	"os"
    10  	"path/filepath"
    11  	"strings"
    12  )
    13  
    14  const Separator = string(filepath.Separator)
    15  
    16  // TarDirToByteBuffer tar归档目录，并写入到ByteBuffer中
    17  func TarDirToByteBuffer(rpath string, compress string) (*bytebufferpool.ByteBuffer, error) {
    18  	b := bytebufferpool.Get()
    19  	defer bytebufferpool.Put(b)
    20  	tw := tar.NewWriter(b)
    21  	defer tw.Close()
    22  	//tar压缩
    23  	err := filepath.Walk(rpath,
    24  		func(path string, info os.FileInfo, err error) error {
    25  			if info.IsDir() {
    26  				//跳过文件夹
    27  				return nil
    28  			}
    29  			header, err := tar.FileInfoHeader(info, "")
    30  			if err != nil {
    31  				return err
    32  			}
    33  			header.Name = strings.TrimPrefix(path, rpath+Separator)
    34  			err = tw.WriteHeader(header)
    35  			if err != nil {
    36  				return err
    37  			}
    38  			filebytes, err := os.ReadFile(path)
    39  			if err != nil {
    40  				return err
    41  			}
    42  			_, err = io.Copy(tw, bytes.NewReader(filebytes))
    43  			if err != nil {
    44  				return err
    45  			}
    46  			return nil
    47  		})
    48  	if err != nil {
    49  		return nil, err
    50  	}
    51  	b2 := bytebufferpool.Get()
    52  	switch compress {
    53  	case "br":
    54  		_, err = fasthttp.WriteBrotliLevel(b2, b.B, fasthttp.CompressBrotliBestCompression)
    55  		if err != nil {
    56  			return nil, err
    57  		}
    58  		return b2, nil
    59  	case "gz":
    60  		_, err = fasthttp.WriteGzipLevel(b2, b.B, fasthttp.CompressBestCompression)
    61  		if err != nil {
    62  			return nil, err
    63  		}
    64  		return b2, nil
    65  	default:
    66  		b2.Write(b.B)
    67  	}
    68  	return b2, nil
    69  }
    70  
    71  // TarDirToFile tar归档目录，并写入到topath中
    72  func TarDirToFile(rpath string, compress string, topath string) error {
    73  	buffer, err := TarDirToByteBuffer(rpath, compress)
    74  	if err != nil {
    75  		return err
    76  	}
    77  	defer bytebufferpool.Put(buffer)
    78  	var suffix string
    79  	switch compress {
    80  	case "br":
    81  		suffix = ".tar.br"
    82  	case "gz":
    83  		suffix = ".tar.gz"
    84  	default:
    85  		suffix = ".tar"
    86  	}
    87  	err = os.WriteFile(topath+suffix, buffer.B, os.ModePerm)
    88  	if err != nil {
    89  		return err
    90  	}
    91  	return nil
    92  }