base-golang/pkg/downloader/controller/controller.go
2024-09-06 15:37:13 +08:00

185 lines
3.7 KiB
Go

package controller
import (
"golang.org/x/net/proxy"
"net"
"net/http"
"net/url"
"os"
"time"
)
type Controller interface {
Touch(name string, size int64) (file *os.File, err error)
Open(name string) (file *os.File, err error)
Write(name string, offset int64, buf []byte) (int, error)
Close(name string) error
ContextDialer() (proxy.Dialer, error)
ContextCookie() http.CookieJar
ContextTimeout() time.Duration
ContextInsecureSkipVerify() bool
ContextProxy() func(*http.Request) (*url.URL, error)
}
type Option func(*option)
type option struct {
CookieJar http.CookieJar
Timeout time.Duration
Dialer proxy.Dialer
InsecureSkipVerify bool
Proxy func(*http.Request) (*url.URL, error)
}
func WithCookie(cookieJar http.CookieJar) Option {
return func(opt *option) {
opt.CookieJar = cookieJar
}
}
func WithTimeout(timeout time.Duration) Option {
return func(opt *option) {
opt.Timeout = timeout
}
}
func WithDialer(dialer proxy.Dialer) Option {
return func(opt *option) {
opt.Dialer = dialer
}
}
func WithProxy(fn func(*http.Request) (*url.URL, error)) Option {
return func(opt *option) {
opt.Proxy = fn
}
}
func WithInsecureSkipVerify(insecure bool) Option {
return func(opt *option) {
opt.InsecureSkipVerify = insecure
}
}
type DefaultController struct {
*option
Files map[string]*os.File
}
func NewController(options ...Option) *DefaultController {
opt := new(option)
for _, f := range options {
f(opt)
}
if opt.Timeout == 0 {
opt.Timeout = time.Second * 30
}
if opt.Dialer == nil {
opt.Dialer = proxy.FromEnvironment()
}
return &DefaultController{
Files: make(map[string]*os.File),
option: opt,
}
}
func (c *DefaultController) Touch(name string, size int64) (file *os.File, err error) {
file, err = os.Create(name)
if size > 0 {
err = os.Truncate(name, size)
if err != nil {
return nil, err
}
}
if err == nil {
c.Files[name] = file
}
return
}
func (c *DefaultController) Open(name string) (file *os.File, err error) {
file, err = os.OpenFile(name, os.O_RDWR, os.ModePerm)
if err == nil {
c.Files[name] = file
}
return
}
func (c *DefaultController) Write(name string, offset int64, buf []byte) (int, error) {
return c.Files[name].WriteAt(buf, offset)
}
func (c *DefaultController) Close(name string) error {
err := c.Files[name].Close()
delete(c.Files, name)
return err
}
func (c *DefaultController) ContextDialer() (proxy.Dialer, error) {
return &DialerWarp{dialer: c.Dialer}, nil
}
func (c *DefaultController) ContextCookie() http.CookieJar {
return c.CookieJar
}
func (c *DefaultController) ContextTimeout() time.Duration {
return c.Timeout
}
func (c *DefaultController) ContextInsecureSkipVerify() bool {
return c.InsecureSkipVerify
}
func (c *DefaultController) ContextProxy() func(*http.Request) (*url.URL, error) {
return c.Proxy
}
type DialerWarp struct {
dialer proxy.Dialer
}
type ConnWarp struct {
conn net.Conn
}
func (c *ConnWarp) Read(b []byte) (n int, err error) {
return c.conn.Read(b)
}
func (c *ConnWarp) Write(b []byte) (n int, err error) {
return c.conn.Write(b)
}
func (c *ConnWarp) Close() error {
return c.conn.Close()
}
func (c *ConnWarp) LocalAddr() net.Addr {
return c.conn.LocalAddr()
}
func (c *ConnWarp) RemoteAddr() net.Addr {
return c.conn.RemoteAddr()
}
func (c *ConnWarp) SetDeadline(t time.Time) error {
return c.conn.SetDeadline(t)
}
func (c *ConnWarp) SetReadDeadline(t time.Time) error {
return c.conn.SetReadDeadline(t)
}
func (c *ConnWarp) SetWriteDeadline(t time.Time) error {
return c.conn.SetWriteDeadline(t)
}
func (d *DialerWarp) Dial(network, addr string) (c net.Conn, err error) {
conn, err := d.dialer.Dial(network, addr)
if err != nil {
return nil, err
}
return &ConnWarp{conn: conn}, nil
}