forked from zmap/zgrab
/
ztls_heartbeat.go
85 lines (70 loc) · 1.46 KB
/
ztls_heartbeat.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
package ztls
import (
"errors"
)
const (
// Record Type
recordTypeHeartbeat recordType = 24
// Extension Number
extensionHeartbeat uint16 = 15
// Heartbeat Mode
heartbeatModePeerAllowed uint8 = 1
heartbeatModePeerNotAllowed uint8 = 2
// Heartbeat Message Types
heartbeatTypeRequest uint8 = 1
heartbeatTypeResponse uint8 = 2
)
var (
HeartbleedError = errors.New("Error after Heartbleed")
)
type Heartbleed struct {
HeartbeatEnabled bool `json:"heartbeat_enabled"`
Vulnerable bool `json:"heartbleed_vulnerable"`
}
type heartbleedMessage struct {
raw []byte
}
func (m *heartbleedMessage) marshal() []byte {
x := make([]byte, 3)
x[0] = 1
x[1] = byte(0x00)
x[2] = byte(0x00)
m.raw = x
return x
}
func (c *Conn) CheckHeartbleed(b []byte) (n int, err error) {
if err = c.Handshake(); err != nil {
return
}
if !c.heartbeat {
return
}
c.in.Lock()
defer c.in.Unlock()
hb := heartbleedMessage{}
hb.marshal()
if _, err = c.writeRecord(recordTypeHeartbeat, hb.raw); err != nil {
return 0, err
}
if err = c.readRecord(recordTypeHeartbeat); err != nil {
return 0, HeartbleedError
}
if c.in.err != nil {
return 0, HeartbleedError
}
n, err = c.input.Read(b)
if c.input.off >= len(c.input.data) {
c.in.freeBlock(c.input)
c.input = nil
}
if n != 0 {
return n, HeartbleedError
}
if err != nil {
return 0, HeartbleedError
}
return 0, HeartbleedError
}
func (c *Conn) GetHeartbleedLog() *Heartbleed {
return c.heartbleedLog
}