forked from scjalliance/smb-http-proxy
-
Notifications
You must be signed in to change notification settings - Fork 0
/
smbconn.go
162 lines (131 loc) · 3.17 KB
/
smbconn.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
package main
import (
"context"
"errors"
"fmt"
"net"
"net/http"
"os"
"strings"
"sync"
"github.com/hirochachacha/go-smb2"
)
// ErrClosed is returned when an operation is attempted on a closed connection.
var ErrClosed = errors.New("smbfs: the connection to the remote file system is closed")
// SMBConn is an individual
type SMBConn struct {
source string
id int
mutex sync.RWMutex
session *smb2.Session
share *smb2.Share
}
// NewConnection establishes a new SMB connection to a remote file system.
func NewConnection(ctx context.Context, conf Config, id int) (*SMBConn, error) {
host, err := ServerAddrFromPath(conf.Source)
if err != nil {
return nil, fmt.Errorf("smbfs: unable to determine source host: %v", err)
}
// Prepare a connetion context to limit the connection time
connCtx, cancel := context.WithTimeout(ctx, conf.ConnTimeout)
defer cancel()
// Prepare a TCP dialer
var tcpd net.Dialer
// Establish a TCP connection via the dialer
conn, err := tcpd.DialContext(connCtx, "tcp", host)
if err != nil {
return nil, err
}
// Prepare an SMB2 dialer
smbd := &smb2.Dialer{
Initiator: &smb2.NTLMInitiator{
User: conf.User,
Password: conf.Password,
Domain: conf.Domain,
},
}
// Establish an SMB2 session via the dialer
session, err := smbd.DialContext(connCtx, conn)
if err != nil {
conn.Close()
return nil, err
}
session = session.WithContext(ctx)
// Mount the desired path
share, err := session.Mount(conf.Source)
if err != nil {
session.Logoff()
return nil, err
}
share = share.WithContext(ctx)
c := &SMBConn{
id: id,
source: conf.Source,
session: session,
share: share,
}
c.log("Connected")
return c, nil
}
// Close disconnects from the SMB server.
func (c *SMBConn) Close() error {
c.mutex.Lock()
defer c.mutex.Unlock()
// Make sure the connection hasn't been closed already
if c.session == nil {
return nil
}
// Execute the logoff process
err := c.session.Logoff()
// Change state to disconnected
c.session = nil
c.share = nil
c.log("Closed")
return err
}
// Open returns a file from the remote file sytem.
func (c *SMBConn) Open(name string) (http.File, error) {
name = strings.TrimPrefix(name, `/`) // Remove leading slash
c.mutex.RLock()
defer c.mutex.RUnlock()
if c.share == nil {
return nil, ErrClosed
}
f, err := c.share.Open(name)
if err != nil {
switch err := err.(type) {
case *os.PathError:
c.log("%s: %v", name, err.Unwrap())
default:
c.log("%v", err)
}
} else {
c.log("%s", name)
}
return f, err
}
// OK returns true if the connection is healthy.
func (c *SMBConn) OK() bool {
c.mutex.RLock()
defer c.mutex.RUnlock()
if c.share == nil {
return false
}
// If we can perform a stat operation on the root of the share it's a
// pretty good indication that the connection is healthy
_, err := c.share.Stat("")
if err != nil {
switch err := err.(type) {
case *os.PathError:
c.log("HEALTH CHECK: %v", err.Unwrap())
default:
c.log("HEALTH CHECK: %v", err)
}
return false
}
c.log("HEALTH OK")
return true
}
func (c *SMBConn) log(format string, v ...interface{}) {
fmt.Printf("[%s][%d]: %s\n", c.source, c.id, fmt.Sprintf(format, v...))
}